Notice how Optimisers.adjust! doesn't mutate train_state (as it should):
julia> using Lux, Optimisers, Random
julia> model = Dense(1=>1)
Dense(1 => 1) # 2 parameters
julia> ps, st = Lux.setup(Random.default_rng(), model)
((weight = Float32[-0.8570953;;], bias = Float32[-0.7004589]), NamedTuple())
julia> opt = Descent(0.1)
Descent(0.1)
julia> train_state = Training.TrainState(model, ps, st, opt)
TrainState(
Dense(1 => 1), # 2 parameters
number of parameters: 2
number of states: 0
optimizer: Descent(0.1)
step: 0
)
julia> Optimisers.adjust!(train_state, 0.234)
TrainState(
Dense(1 => 1), # 2 parameters
number of parameters: 2
number of states: 0
optimizer: Descent(0.234)
step: 0
)
julia> train_state.optimizer # Should return Descent(0.234)
Descent(0.1)
I am using
Status `/tmp/jl_60I9SX/Project.toml`
[b2108857] Lux v1.31.4
[3bd65402] Optimisers v0.4.7
[9a3f8284] Random v1.11.0
Notice how
Optimisers.adjust!doesn't mutatetrain_state(as it should):I am using