Update llama.py

2025-02-17 00:49:35 -08:00 · 2025-02-17 00:49:35 -08:00 · a5fcf2dfc6
commit a5fcf2dfc6
parent a9a4b5e436
1 changed files with 1 additions and 0 deletions
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@ -700,6 +700,7 @@ def LlamaModel_fast_forward(
        elif inputs_requires_grad:
            inputs_embeds.requires_grad_(False)
        pass
+        attention_mask = attention_mask[:,:self.max_seq_length] # Must resize!
        inputs_embeds *= attention_mask.unsqueeze(0).transpose(0, 1).transpose(1, 2)
        if inputs_requires_grad: inputs_embeds.requires_grad_(True)
    pass