From 94aa679684440adac244a6e0b58e77b664dab448 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fabian=20Fr=C3=B6hlich?= <fabian@schaluck.com>
Date: Sat, 16 Nov 2024 23:12:33 +0000
Subject: [PATCH] Update petab.py

---
 python/sdist/amici/jax/petab.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/python/sdist/amici/jax/petab.py b/python/sdist/amici/jax/petab.py
index deb1d12d92..c929217e55 100644
--- a/python/sdist/amici/jax/petab.py
+++ b/python/sdist/amici/jax/petab.py
@@ -23,7 +23,7 @@
     ParameterMappingForCondition,
     create_parameter_mapping,
 )
-from amici.jax.model import JAXModel, simulate_condition
+from amici.jax.model import JAXModel
 
 
 def jax_unscale(
@@ -239,10 +239,9 @@ def run_simulation(
             if len(simulation_condition) > 1
             else jnp.array([])
         )
-        return simulate_condition(
+        return self.model.simulate_condition(
             p,
             p_preeq,
-            self.model,
             jax.lax.stop_gradient(jnp.array(ts_preeq)),
             jax.lax.stop_gradient(jnp.array(ts_dyn)),
             jax.lax.stop_gradient(jnp.array(ts_posteq)),