diff --git a/src/maxtext/layers/attention_mla.py b/src/maxtext/layers/attention_mla.py
index d5f79a1a4b..e81b0a9b73 100644
--- a/src/maxtext/layers/attention_mla.py
+++ b/src/maxtext/layers/attention_mla.py
@@ -1077,6 +1077,8 @@ def calculate_indexer_loss(
 
     # Aggregate heads: [b, h, t, s] -> [b, t, s]
     attention_probs = jnp.sum(attention_probs, axis=1)
+    # Force materialization and prevent fusion across this point to reuse the intermediate tensor
+    attention_probs = jax.lax.optimization_barrier(attention_probs)
     # L1 normalize aggregated target distribution
     attention_probs = attention_probs / (jnp.sum(attention_probs, axis=-1, keepdims=True) + EPS)