From 56fbfb29030abd21a18eb0db3372c02b73535b96 Mon Sep 17 00:00:00 2001 From: Yang Zhao Date: Mon, 6 Nov 2023 19:31:09 +0800 Subject: [PATCH] fix_bug_in_optimizer_api_md --- docs/flip/1009-optimizer-api.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/flip/1009-optimizer-api.md b/docs/flip/1009-optimizer-api.md index f658db5cc..ec8d15737 100644 --- 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