be excplicit with kwargs

axolotl-ai-cloud · Sep 26, 2023 · 367d8ec · 367d8ec
1 parent 68b57b5
commit 367d8ec
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -530,11 +530,11 @@ def custom_forward(*inputs):
             layer_outputs = torch.utils.checkpoint.checkpoint(
                 create_custom_forward(decoder_layer),
                 hidden_states,
-                attention_mask,
-                position_ids,
-                None,
-                cu_seqlens,
-                max_seqlen,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                use_cache=None,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=max_seqlen,
             )
         else:
             layer_outputs = decoder_layer(