fix profile_nn_jax api change

ACEsuit · Mar 11, 2023 · 09a999e · 09a999e
1 parent fe778e7
commit 09a999e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/mace_jax/tools/gin_functions.py b/mace_jax/tools/gin_functions.py
@@ -33,7 +33,7 @@ def flags(
     if profile:
         import profile_nn_jax
 
-        profile_nn_jax.enable(timing=True, statistics=True)
+        profile_nn_jax.enable()
     return seed