Optimization.jl and Lux.jl 1.10.0 Compatability

I was revisiting code I wrote in Julia 1.9.3 which closely mimics the code written by Chris here. While it worked fine in 1.9.3, when I updated my packages and installation to 1.10.0. I started getting the following error

LoadError: type OptimizationState has no field layer_1

Seems to link back to the way Optimization.solve links back to the Lux NN. Here is a snippet of the code in my file.

rbf(x) = exp.(-(x .^ 2))
# Multilayer FeedForward
U = Lux.Chain(Lux.Dense(1, 5, rbf), Lux.Dense(5, 5, rbf), Lux.Dense(5, 5, tanh),
    Lux.Dense(5, 1))
# Get the initial parameters and state variables of the model
p, st = Lux.setup(rng, U);

# Define the hybrid model
function ude_dynamics!(du, u, p, t, p_true)
    û = U(u, p, st)[1]
    du[1] = p_true[1] * û[1]

# Closure with the known parameter
nn_dynamics!(du, u, p, t) = ude_dynamics!(du, u, p, t, p_)

# Define the UDE problem
prob_nn = ODEProblem(nn_dynamics!, u0, tspan, p)

which, after some function definitions, continues to

#Define the optimization function with automatic differentiation
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p)->loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentVector{Float64}(p))

res1 = Optimization.solve(optprob, ADAM(), callback = callback, maxiters = 1000)

The error is thrown at res1 = Optimization.solve(optprob, ADAM(), callback = callback, maxiters = 1000). Any help is appreciated.

You dropped the important part.