diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py
index 3d0560e140..7e68c5c34d 100644
--- a/python/sdist/amici/jax/petab.py
+++ b/python/sdist/amici/jax/petab.py
@@ -644,7 +644,7 @@ def petab_simulate(
     for sc, ys in y.items():
         obs = [
             problem.model.observable_ids[io]
-            for io in problem._measurements[sc][4]
+            for io in problem._measurements[sc][3]
         ]
         t = jnp.concat(problem._measurements[sc][:2])
         df_sc = pd.DataFrame(