diff --git a/mjx/training_apg.ipynb b/mjx/training_apg.ipynb
index 66a9b8d6..b99ebc45 100644
--- a/mjx/training_apg.ipynb
+++ b/mjx/training_apg.ipynb
@@ -633,6 +633,7 @@
     "                             learning_rate=1e-4,\n",
     "                             num_eval_envs=64,\n",
     "                             num_evals=10 + 1,\n",
+    "                             use_float64=True,\n",
     "                             normalize_observations=True,\n",
     "                             network_factory=make_networks_factory)"
    ]
@@ -1223,6 +1224,7 @@
     "                             schedule_decay=0.995,\n",
     "                             num_eval_envs=64,\n",
     "                             num_evals=10 + 1,\n",
+    "                             use_float64=True,\n",
     "                             normalize_observations=True,\n",
     "                             network_factory=make_networks_factory)\n",
     "\n",