From 5d628d9397926d44b73d843d6665f91bf2e63872 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 12 Feb 2025 03:32:12 -0800 Subject: [PATCH] Update llama.py --- unsloth/models/llama.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index fb05e052dd..e03f73301d 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1154,7 +1154,13 @@ def CausalLM_fast_forward(fast_forward_inference): output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output - print(loss, logits) + print(CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + )) return CausalLMOutputWithPast( loss=loss, logits=logits,