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",