Zygote error: 'try/catch is not supported'

Hello,

I would like to use Zygote with ITensor to get gradients via autodiff. I have tried to implement a simple example where my loss function measures the inner product between two MPOs, however I get an error message Compiling Tuple{typeof(lognorm), MPO}: try/catch is not supported. It seems like this is due to the ‘inner’ function, but I am not sure why as in the examples folder you use ‘inner’ in the loss function without this error flagging up.

Please can you help resolve this, many thanks.

using ITensors
using OptimKit
using Random
using Zygote

X = [0 1; 1 0]
XX = kron(X, X)

n = 2

s = siteinds("Qubit", n)

function loss(θ⃗)
    
    mpo_init = MPO(s, "Id")
    mpo1 = apply(op(exp(-im*θ⃗[1]*XX), s[1], s[2]), mpo_init)
    mpo2 = apply(op(exp(-im*1*XX), s[1], s[2]), mpo_init)    
    return -abs(inner(mpo1, mpo2))^2

end

θ⃗ = 2π * rand(1)

@show loss(θ⃗)
@show loss'(θ⃗)