diff --git a/pipelines/lumina_dimmo/lumina_dimoo.py b/pipelines/lumina_dimmo/lumina_dimoo.py
index d16ca17ef..9a4b90898 100644
--- a/pipelines/lumina_dimmo/lumina_dimoo.py
+++ b/pipelines/lumina_dimmo/lumina_dimoo.py
@@ -874,10 +874,10 @@ class LLaDASequentialBlock(LLaDABlock):
 
         if self._activation_checkpoint_fn is not None:
             att, cache = self._activation_checkpoint_fn(  # type: ignore
-                self.attention, q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache
+                self.attention, q, k, v, attention_bias, layer_past=layer_past
             )
         else:
-            att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache)
+            att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past)
 
         x = x + self.dropout(att)
 
@@ -967,7 +967,7 @@ class LLaDALlamaBlock(LLaDABlock):
 
         if self._activation_checkpoint_fn is not None:
             att, cache = self._activation_checkpoint_fn(  # type: ignore
-                self.attention, q, k, v, attention_bias, layer_past=layer_past, use_cache=use_cache
+                self.attention, q, k, v, attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask
             )
         else:
             att, cache = self.attention(q, k, v, attention_bias, layer_past=layer_past, to_compute_mask=to_compute_mask)