diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index e9ac9bfaf5..f0dcb13a94 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -232,7 +232,6 @@ def LlamaAttention_fast_forward(
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
     
     bsz, q_len, _ = hidden_states.size()
-    print(hidden_states.size())
 
     # Check for inference
     if past_key_value is not None and q_len == 1 and bsz == 1:
diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index baf3695b22..6d14c98d5d 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -49,7 +49,7 @@ def MistralAttention_fast_forward(
     bsz, q_len, _ = hidden_states.size()
 
     # Check for inference
-    if past_key_value is not None and q_len == 1:
+    if past_key_value is not None and q_len == 1 and bsz == 1:
         A, past_key_value = LlamaAttention_fast_forward_inference(
             self,
             hidden_states,