fix convlution module norm dim (#2684)

wenet-e2e · Feb 7, 2025 · ead4d14 · ead4d14
1 parent 0228877
commit ead4d14
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/wenet/transformer/convolution.py b/wenet/transformer/convolution.py
@@ -76,11 +76,15 @@ def __init__(
         assert norm in ['batch_norm', 'layer_norm', 'rms_norm']
         if norm == "batch_norm":
             self.use_layer_norm = False
-            self.norm = WENET_NORM_CLASSES['batch_norm'](channels,
+            self.norm = WENET_NORM_CLASSES['batch_norm'](conv_inner_factor *
+                                                         channels // 2,
                                                          eps=norm_eps)
         else:
             self.use_layer_norm = True
-            self.norm = WENET_NORM_CLASSES[norm](channels, eps=norm_eps)
+            self.norm = WENET_NORM_CLASSES[norm](conv_inner_factor *
+                                                 channels // 2,
+                                                 eps=norm_eps)
+
 
         self.pointwise_conv2 = nn.Conv1d(
             conv_inner_factor * channels // 2,