Fix jit

RobertTLange · Oct 12, 2024 · cdc622b · cdc622b
1 parent 04f9ad4
commit cdc622b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/evosax/strategy.py b/evosax/strategy.py
@@ -85,7 +85,7 @@ def initialize(
             state = self.set_mean(state, init_mean)
         return state
 
-    # @partial(jax.jit, static_argnums=(0,))
+    @partial(jax.jit, static_argnums=(0,))
     def ask(
         self,
         rng: chex.PRNGKey,
@@ -109,7 +109,7 @@ def ask(
             x_out = x_clipped
         return x_out, state
 
-    # @partial(jax.jit, static_argnums=(0,))
+    @partial(jax.jit, static_argnums=(0,))
     def tell(
         self,
         x: Union[chex.Array, chex.ArrayTree],