renmae apply_gradients

LuxDL · May 13, 2024 · f692da2 · f692da2
1 parent 35c7648
commit f692da2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/test/contrib/training_tests.jl b/test/contrib/training_tests.jl
@@ -88,7 +88,7 @@ end
  for epoch in 1:100, (x, y) in dataset_
  grads, loss, _, tstate = Lux.Experimental.compute_gradients(
  ad, mse, (x, y), tstate)
- tstate = Lux.Experimental.apply_gradients(tstate, grads, true)
+ tstate = Lux.Experimental.apply_gradients!(tstate, grads)
  end
 
  final_loss = first(mse(model, tstate.parameters, tstate.states, dataset_[1]))