Update predict_test.py

google · Aug 24, 2023 · db53e46 · db53e46
1 parent 0f9bb34
commit db53e46
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/predict_test.py b/tests/predict_test.py
@@ -300,7 +300,7 @@ def testNTKMeanCovPrediction(
     self.assertGreater(np.min(np.linalg.eigh(cov_train_inf)[0]), -1e-8)
 
     _kernel_fn = nt.empirical_kernel_fn(f)
-    kernel_fn = jit(lambda x1, x2, params: _kernel_fn(x1, x2, 'ntk', params))
+    kernel_fn = lambda x1, x2, params: _kernel_fn(x1, x2, 'ntk', params)
 
     def predict_empirical(key):
       _, params = init_fn(key, train_shape)