DeepONet with additional layer exit error

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

Start by running the Julia reproduction in the issue with the additional DeepONet layer and Zygote.gradient, capturing the complete exit error. Compare it with the version without additional and isolate whether the failure is in DeepONet construction, evaluation, or differentiation. Done means identifying the cause and confirming the corrected behavior with the reproduction.