diff --git a/mjx/tutorial.ipynb b/mjx/tutorial.ipynb
index d462fa5e..e9b983de 100644
--- a/mjx/tutorial.ipynb
+++ b/mjx/tutorial.ipynb
@@ -186,6 +186,7 @@
     "from brax import envs\n",
     "from brax import math\n",
     "from brax.base import Base, Motion, Transform\n",
+    "from brax.base import State as PipelineState\n",
     "from brax.envs.base import Env, PipelineEnv, State\n",
     "from brax.mjx.base import State as MjxState\n",
     "from brax.training.agents.ppo import train as ppo\n",
@@ -1636,6 +1637,7 @@
    "cell_type": "code",
    "execution_count": 0,
    "metadata": {
+    "cellView": "form",
     "id": "AuU-9nUquEu4"
    },
    "outputs": [],