fix notebook

AMICI-dev · Dec 7, 2024 · 17e2ef2 · 17e2ef2
1 parent 01cd6bb
commit 17e2ef2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb
@@ -352,6 +352,7 @@
    "source": [
     "import jax.numpy as jnp\n",
     "import diffrax\n",
+    "from amici.jax import ReturnValue\n",
     "\n",
     "# Define the simulation condition\n",
     "simulation_condition = (\"model1_data1\",)\n",
@@ -380,7 +381,7 @@
     "        controller=diffrax.PIDController(atol=1e-8, rtol=1e-8),\n",
     "        max_steps=2**10,\n",
     "        adjoint=diffrax.DirectAdjoint(),\n",
-    "        ret=\"y\",  # Return observables\n",
+    "        ret=ReturnValue.y,  # Return observables\n",
     "    )[0]\n",
     "\n",
     "\n",