fix hessian

AMICI-dev · Oct 25, 2024 · 8b9c10a · 8b9c10a
1 parent 480b75a
commit 8b9c10a
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/python/sdist/amici/jax.py b/python/sdist/amici/jax.py
@@ -258,7 +258,15 @@ def s2run(
         )(ts, ts_dyn, p, k, k_preeq, my, pscale, dynamic=dynamic)
 
         s2llh = jax.hessian(self._run, 2, True)(
-            ts, ts_dyn, p, k, k_preeq, my, pscale, dynamic=dynamic
+            ts,
+            ts_dyn,
+            p,
+            k,
+            k_preeq,
+            my,
+            pscale,
+            checkpointed=False,
+            dynamic=dynamic,
         )
 
         return llh, sllh, s2llh, (x, obs, stats)