Update llama.py

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,