Remove jitting from ntk calculation (#114)

zincware · Apr 8, 2024 · a60a80b · a60a80b
1 parent cc82976
commit a60a80b
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/znnl/models/jax_model.py b/znnl/models/jax_model.py
@@ -127,7 +127,6 @@ def __init__(
             batch_size=ntk_batch_size,
             store_on_device=store_on_device,
         )
-        self.empirical_ntk_jit = jax.jit(self.empirical_ntk)
         self.apply_jit = jax.jit(self.apply)
 
     def init_model(
@@ -249,7 +248,7 @@ def compute_ntk(
         """
         if x_j is None:
             x_j = x_i
-        empirical_ntk = self.empirical_ntk_jit(
+        empirical_ntk = self.empirical_ntk(
             x_i,
             x_j,
             {