Skip to content

Optimisers.adjust! doesn't mutate TrainState #1722

Description

@Sleort

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

Metadata

Metadata

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions