fix checkpoint forward kwargs

axolotl-ai-cloud · Sep 26, 2023 · bb87204 · bb87204
1 parent 9dbdcec
commit bb87204
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/axolotl/monkeypatch/llama_attn_hijack_flash.py b/src/axolotl/monkeypatch/llama_attn_hijack_flash.py
@@ -522,7 +522,7 @@ def custom_forward(*inputs):
                         *inputs,
                         past_key_value,  # pylint: disable=(cell-var-from-loop)
                         output_attentions,
-                        attention_mask=attention_mask,
+                        padding_mask=padding_mask,
                     )
 
                 return custom_forward