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,