From fbf6cc747ef52812f5593a2daf96e55b5c200514 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Thu, 18 Jul 2024 13:33:13 -0700
Subject: [PATCH] Update mistral.py

---
 unsloth/models/mistral.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/unsloth/models/mistral.py b/unsloth/models/mistral.py
index d7376d9521..b2531056a0 100644
--- a/unsloth/models/mistral.py
+++ b/unsloth/models/mistral.py
@@ -161,7 +161,7 @@ def MistralAttention_fast_forward(
         A = A.transpose(1, 2).contiguous()
     pass
     
-    attn_output = A.reshape(bsz, q_len, self.hidden_size)
+    attn_output = A.reshape(bsz, q_len, n_heads*head_dim)
     attn_output = self.apply_o(self, attn_output)
     attn_weights = None
     return attn_output, attn_weights, past_key_value