Upload modeling.py with huggingface_hub

modeling.py

@@ -284,8 +284,10 @@ class GroupedQueryAttention(nn.Module):
         attention = self.dropout(attention)
         context = torch.matmul(attention, v)
 
-        # Fix context reshaping - always use d_model for output projection
-        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
+        # Fix context reshaping - properly handle grouped query attention output
+        # context shape: [batch_size, num_heads, seq_len, d_k]
+        # We need to reshape it to [batch_size, seq_len, d_model] where d_model = num_heads * d_k
+        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_heads * self.d_k)
         
         output = self.W_o(context)