perf: try AoT compile for jax

LuxDL · Nov 5, 2024 · 8d2209e · 8d2209e
1 parent 0537f57
commit 8d2209e
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/perf/resnet/main.py b/perf/resnet/main.py
@@ -163,7 +163,9 @@ def __call__(self, x, train: bool = True):
 
         print(f"Param count: {param_count}")
 
-        apply_fn_compiled = jax.jit(partial(model.apply, train=False))
+        apply_fn_compiled = (
+            jax.jit(partial(model.apply, train=False)).lower(params, x).compile()
+        )
 
         best_timing = np.inf
         for i in range(100):