Update modeling_auristream.py

modeling_auristream.py

@@ -552,8 +552,8 @@ class CausalSelfAttention(nn.Module):
 
         if self.rotary is not None:
             cos, sin = self.rotary(q)
-            q = apply_rotary_emb(q.transpose(1, 2), cos, sin).transpose(1, 2)
-            k = apply_rotary_emb(k.transpose(1, 2), cos, sin).transpose(1, 2)
+            q = apply_rotary_emb(q, cos, sin)
+            k = apply_rotary_emb(k, cos, sin)
 
         # manual implementation of attention
         att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))