From d6e85d77310cf77d33e672592975aa90c2bd7afe Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Tue, 23 Jan 2024 19:17:42 +1100 Subject: [PATCH] q_len == 1 --- unsloth/models/llama.py | 1 - unsloth/models/mistral.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) 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,