This code causes an error exiting Julia during Zygote.gradient . Any ideas why? Without additional layer works fine.
try
deeponet = NeuralOperators.DeepONet(
Chain(
Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast), Dense(10 => 10)),
Chain(Dense(1 => 10, Lux.tanh_fast), Dense(10 => 10, Lux.tanh_fast),
Dense(10 => 10, Lux.tanh_fast)),
additional = Chain(Dense(10 => 10, Lux.tanh_fast), Dense(10 => 1)))
u = rand(1, 50)
ur = rand(1, 50)
v = rand(1, 40, 1)
θ, st = Lux.setup(Random.default_rng(), deeponet)
ff =(θ) -> deeponet((u, v), θ, st)[1] .- deeponet((ur, v), θ, st)[1]
using Zygote
Zygote.gradient((θ) -> sum(ff(θ)), θ)
catch err
throw(err)
end
This code causes an error exiting Julia during
Zygote.gradient. Any ideas why? Withoutadditionallayer works fine.