fix norm

wenet-e2e · Aug 7, 2024 · f599bf8 · f599bf8
1 parent d22745f
commit f599bf8
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/wenet/transformer/attention.py b/wenet/transformer/attention.py
@@ -668,6 +668,9 @@ def forward(
         q = WENET_APPLY_ROTARY_EMB[self.style](q, pos_emb)
         k = WENET_APPLY_ROTARY_EMB[self.style](k, pos_emb)
 
+        if self.qk_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
         k, v, new_cache = self._update_kv_and_cache(k,
                                                     v,
                                                     cache,