diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index 6314245a..751a150a 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -1513,8 +1513,8 @@
    "source": [
     "# grab the latest checkpoint from the flat terrain joystick policy\n",
     "latest_ckpts = list(ckpt_path.glob('*'))\n",
-    "latest_ckpts.sort()\n",
-    "latest_ckpt = latest_ckpts[0]\n",
+    "latest_ckpts.sort(key=lambda x: int(x.as_posix().split('/')[-1]))\n",
+    "latest_ckpt = latest_ckpts[-1]\n",
     "\n",
     "train_fn = functools.partial(\n",
     "      ppo.train, num_timesteps=40_000_000, num_evals=5,\n",