From 17e2ef21555c28374171d6b2699411a09309a7f7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= <fabian@schaluck.com>
Date: Sat, 7 Dec 2024 20:09:19 +0000
Subject: [PATCH] fix notebook

---
 python/examples/example_jax_petab/ExampleJaxPEtab.ipynb | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb b/python/examples/example_jax_petab/ExampleJaxPEtab.ipynb
index 121c082eaa..1310091f4c 100644
--- 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",