Update modeling_gemmoe.py

modeling_gemmoe.py

@@ -243,8 +243,9 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
     Returns:
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
-    cos = cos.unsqueeze(unsqueeze_dim)
-    sin = sin.unsqueeze(unsqueeze_dim)
+    seq_len, dim = q.shape[-2], q.shape[-1]
+    cos = cos[:seq_len].view(1, 1, seq_len, dim)
+    sin = sin[:seq_len].view(1, 1, seq_len, dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed