From 1d7b19eac51348625cb0e5b76f8785047f9cf124 Mon Sep 17 00:00:00 2001 From: rootkiller6788 Date: Thu, 20 Aug 2026 02:50:15 +0800 Subject: [PATCH] Fix typo in Part1 tutorial notebook The MomentumOptimizer.update example populates the `momentums` field of MomentumOptState from `output_params` instead of `output_momentums`, so the returned state carries the updated parameters as its momentum. Use `output_momentums` to match the intended momentum update. --- docs/notebooks/Part1_Introduction.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/notebooks/Part1_Introduction.ipynb b/docs/notebooks/Part1_Introduction.ipynb index a7528708..6d390fad 100644 --- a/docs/notebooks/Part1_Introduction.ipynb +++ b/docs/notebooks/Part1_Introduction.ipynb @@ -776,7 +776,7 @@ " params=jax.tree_util.tree_unflatten(struct, output_params),\n", " model_state=model_state,\n", " iteration=opt_state.iteration + 1,\n", - " momentums=jax.tree_util.tree_unflatten(struct, output_params),\n", + " momentums=jax.tree_util.tree_unflatten(struct, output_momentums),\n", " )\n", "\n", "\n",