Merge pull request #3462 from zhaoyang-0204:my_branch

PiperOrigin-RevId: 582824459
google · Nov 15, 2023 · 78b0c8c · 78b0c8c
2 parents 79d925f + 56fbfb2
commit 78b0c8c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docs/flip/1009-optimizer-api.md b/docs/flip/1009-optimizer-api.md
@@ -105,7 +105,7 @@ def train_step(opt_state, variables, inputs, labels, apply_fn, tx_update_fn):
       params)
   updates, new_opt_state = tx_update_fn(grads, opt_state, params)
   new_params = optax.apply_updates(params, updates)
-  new_variables = {**variables, **new_model_state, 'params': params}
+  new_variables = {**variables, **new_model_state, 'params': new_params}
   return new_opt_state, new_variables, loss