Fix crashing Github tests

PiperOrigin-RevId: 559848201
google · Aug 24, 2023 · 91846d8 · 91846d8
1 parent b3b728e
commit 91846d8
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/tests/predict_test.py b/tests/predict_test.py
@@ -300,7 +300,8 @@ 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))
+    # TODO(romann): figure out the slow compile on Ubuntu 22.04 CPU Python 3.9
+    kernel_fn = lambda x1, x2, params: _kernel_fn(x1, x2, 'ntk', params)
 
     def predict_empirical(key):
       _, params = init_fn(key, train_shape)