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]
end
# 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.