remove static

google · May 18, 2024 · 53e5740 · 53e5740
1 parent 2ca0d24
commit 53e5740
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/MaxText/maxengine.py b/MaxText/maxengine.py
@@ -179,7 +179,7 @@ def prefill(
  "generated_tokens": generated_tokens,
  }
 
- @functools.partial(jax.jit, static_argnums=(0,), donate_argnums=(2,))
+ @functools.partial(jax.jit, donate_argnums=(2,))
  def generate(self, params: Params, decode_state: DecodeState) -> Tuple[DecodeState, engine_api.ResultTokens]:
  """Run one generate step"""
  previous_logits = decode_state["logits"]