From 9ac3754fcf79c1285ea118c849e18e63108af831 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Jun 2025 01:57:15 -0700
Subject: [PATCH] Update llama.py

---
 unsloth/models/llama.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index 9db8abdd43..0cdd51733b 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -1130,7 +1130,9 @@ def CausalLM_fast_forward(fast_forward_inference):
         # Output last hidden states without logits if asked
         if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1":
             if num_logits_to_keep != 0:
+                print("llama 1133 hidden_states", hidden_states.shape)
                 hidden_states = hidden_states[:, -num_logits_to_keep:, :]
+                print("llama 1133 hidden_states", hidden_states.shape)
             return CausalLMOutputWithPast(
                 loss = None,
                 logits = hidden_states,