remove jit

zincware · May 7, 2024 · c5e1170 · c5e1170
1 parent afb849f
commit c5e1170
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/znnl/models/jax_model.py b/znnl/models/jax_model.py
@@ -122,7 +122,7 @@ def __init__(
 
         # Prepare NTK calculation
         self.empirical_ntk = nt.batch(
-            nt.empirical_ntk_fn(f=jax.jit(self._ntk_apply_fn), trace_axes=trace_axes),
+            nt.empirical_ntk_fn(f=self._ntk_apply_fn, trace_axes=trace_axes),
             batch_size=ntk_batch_size,
             store_on_device=store_on_device,
         )