diff --git a/flaxdiff/samplers/common.py b/flaxdiff/samplers/common.py
index 186a135..e1dfe84 100644
--- a/flaxdiff/samplers/common.py
+++ b/flaxdiff/samplers/common.py
@@ -133,7 +133,7 @@ def generate_images(self,
             
         params = params if params is not None else self.params
         
-        @jax.jit
+        # @jax.jit
         def sample_model_fn(x_t, t, *additional_inputs):
             return self.sample_model(params, x_t, t, *additional_inputs)
 
diff --git a/pyproject.toml b/pyproject.toml
index 606d470..6f73252 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "flaxdiff"
-version = "0.1.37.6"
+version = "0.1.37.7"
 description = "A versatile and easy to understand Diffusion library"
 readme = "README.md"
 authors = [