Fix positional arguments to optimizer.stateless_apply (#1732)

guides/writing_a_custom_training_loop_in_jax.py

@@ -215,7 +215,7 @@ def train_step(state, data):
         trainable_variables, non_trainable_variables, x, y
     )
     trainable_variables, optimizer_variables = optimizer.stateless_apply(
-        grads, trainable_variables, optimizer_variables
+        optimizer_variables, grads, trainable_variables
     )
     # Return updated state
     return loss, (