Update modeling_quiet.py

modeling_quiet.py

@@ -158,15 +158,13 @@ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
     Returns:
         `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
     """
-    seq_len = q.size(-2)
-    cos = cos[position_ids, :, :seq_len, :]
-    sin = sin[position_ids, :, :seq_len, :]
+    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
     q_embed = (q * cos) + (rotate_half(q) * sin)
     k_embed = (k * cos) + (rotate_half(k) * sin)
     return q_embed, k_embed
 
 
-
 class QuietMLP(nn.Module):
     def __init__(self, config):
         super().__init__()