From e9fda7ab1dcf9df777c478de1f7ef4fad94f16c9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 13 Feb 2025 00:10:44 -0800 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index af144f01a1..0b567b023a 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -707,9 +707,9 @@ def LlamaModel_fast_forward( # Ignore attention_mask if attention_mask is None: padding_mask = None - # elif self.training: - # attention_mask = None - # padding_mask = None + elif self.training: + attention_mask = None + padding_mask = None else: # if 0 in attention_mask: # padding_mask = attention_mask @@ -1161,7 +1161,7 @@ def CausalLM_fast_forward(fast_forward_inference): if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output - print("***", logits.dtype, logits.shape) + return CausalLMOutputWithPast( loss=loss, logits=logits,