diff --git a/tunix/models/gemma4/attention.py b/tunix/models/gemma4/attention.py
index 42a1d87c0..cc6b36f43 100644
--- a/tunix/models/gemma4/attention.py
+++ b/tunix/models/gemma4/attention.py
@@ -879,7 +879,7 @@ def __call__(
 
       def _checkpointed_block(state, *args, **kwargs):
         module = nnx.merge(graphdef, state)
-        return module.block(*args, **kwargs)
+        return module.block(*args, **kwargs, force_eager=force_eager)
 
       return jax.checkpoint(_checkpointed_block)(
           state,
@@ -889,7 +889,6 @@ def _checkpointed_block(state, *args, **kwargs):
           attn_mask,
           kv_shared_cache,
           segment_ids,
-          force_eager,
       )
     else:
       return self.block(