From 5c927e4930e2d52661b440c4ca1cc999a6cdb2a3 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 21 Jan 2024 19:01:28 +1100 Subject: [PATCH] Update mistral.py --- unsloth/models/mistral.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py index 2c85129010..68add4b239 100644 --- a/unsloth/models/mistral.py +++ b/unsloth/models/mistral.py @@ -47,10 +47,9 @@ def MistralAttention_fast_forward( ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() - Q, K, V = self.apply_qkv(self, hidden_states) # Check for inference - if use_cache and past_key_value is not None and q_len == 1: + if past_key_value is not None and q_len == 1: A, past_key_value = LlamaAttention_fast_forward_inference( self, hidden_states, @@ -66,6 +65,7 @@ def MistralAttention_fast_forward( head_dim = self.head_dim assert(n_kv_heads * n_groups == n_heads) + Q, K, V = self.apply_qkv(self, hidden_states) Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2) K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2) V = V.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2)