fix: rope dim mismatch

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

rotary.py

@@ -31,6 +31,10 @@ def apply_rotary_emb_torch(x, cos, sin, interleaved=False):
     """
     ro_dim = cos.shape[-1] * 2
     assert ro_dim <= x.shape[-1]
+    cos, sin = (
+        cos[:x.shape[1]],
+        sin[:x.shape[1]],
+    )
     cos = repeat(cos, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
     sin = repeat(sin, "... d -> ... 1 (2 d)" if not interleaved else "... d -> ... 1 (d 2)")
     return torch.cat(