Merge pull request #4857 from nenuadrian:patch-1

PiperOrigin-RevId: 797931619

Author: Flax Authors

examples/nnx_toy_examples/05_vae.py

@@ -138,7 +138,7 @@ def loss_fn(model: VAE):
     return loss
 
   loss, grads = nnx.value_and_grad(loss_fn)(model)
-  optimizer.update(model, grads)
+  optimizer.update(grads)
 
   return loss