diff --git a/notebooks/03-MultiGPU_PM_Halo.ipynb b/notebooks/03-MultiGPU_PM_Halo.ipynb
index 025b715..be0bfb8 100644
--- a/notebooks/03-MultiGPU_PM_Halo.ipynb
+++ b/notebooks/03-MultiGPU_PM_Halo.ipynb
@@ -177,7 +177,7 @@
     "\n",
     "    # Evolve the simulation forward\n",
     "    ode_fn = ODETerm(\n",
-    "        make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False , sharding=sharding))\n",
+    "        make_diffrax_ode(mesh_shape, paint_absolute_pos=False , sharding=sharding))\n",
     "    solver = Tsit5()\n",
     "\n",
     "    stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\n",
@@ -407,7 +407,7 @@
     "\n",
     "    # Evolve the simulation forward\n",
     "    ode_fn = ODETerm(\n",
-    "        make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False , sharding=sharding , halo_size=halo_size))\n",
+    "        make_diffrax_ode(mesh_shape, paint_absolute_pos=False , sharding=sharding , halo_size=halo_size))\n",
     "    solver = Tsit5()\n",
     "\n",
     "    stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\n",