Update modeling_gpt2_mq.py

modeling_gpt2_mq.py

@@ -244,7 +244,7 @@ class GPT2MQAttention(nn.Module):
             attention_mask = encoder_attention_mask
         else:
             query = self.q_attn(hidden_states)
-            key, value = self.kv_attn(hidden_states).split(self.split_size, dim=2)
+            key, value = self.kv_attn(hidden_states).split(self.head_dim, dim=2)
 
 
         batch_size, seq_length = query.shape[:2]