FIX

Author: KakaruHayate

modules/commons/common_layers.py

@@ -275,16 +275,16 @@ def forward(self, x, key_padding_mask=None):
         # Project inputs to Q, K, V
         Q, K, V = torch.split(self.in_proj(x), self.embed_dim, dim=-1)
 
-        # Query-Key Normalization
-        if self.use_qk_norm:
-            Q = self.q_norm(Q)
-            K = self.k_norm(K)
-        
         # Reshape Q, K, V for multi-head attention
         Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
         K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
         V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
 
+        # Query-Key Normalization
+        if self.use_qk_norm:
+            Q = self.q_norm(Q)
+            K = self.k_norm(K)
+        
         # Apply RoPE
         if self.rotary_embed is not None:
             Q = self.rotary_embed.rotate_queries_or_keys(Q)