Hi,

I’m working on hybrid systems with the SciML framework, approximating one of the derivatives of an ODE system with a neural network. When I attempt, using Lux.freeze, to freeze some of the network parameters that I don’t want to be optimized (for example the weights in the snippet that I’ve attached), I get the following error during optimization:

ERROR: MethodError: no method matching _merge(::ReverseDiff.TrackedArray{Float64, Float64, 1, ComponentVector{Float64, SubArray{Float64, 1, Vector{Float64}, Tuple{UnitRange{Int64}}, true}, Tuple{Axis{(weight = ViewAxis(1:10, ShapedA …

snippet_fail_frozen_layer.jl (2.3 KB)

Could you kindly point out what I am doing wrong?

Thank you in advance, Stefano

I am using Julia v"1.8.5" with the following packages

ComponentArrays v0.13.7

DifferentialEquations v7.7.0

Lux v0.4.36

Optimization v3.11.2

OptimizationOptimisers v0.1.1

SciMLSensitivity v7.20.0

StableRNGs v1.0.0