conformer variant fix

mlcommons · Nov 16, 2023 · 88bfde9 · 88bfde9
1 parent 86bf643
commit 88bfde9
Show file tree

Hide file tree

Showing 4 changed files with 5 additions and 5 deletions.
diff --git a/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/models.py b/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/models.py
@@ -60,7 +60,7 @@ class ConformerConfig:
   use_specaug: bool = True
   attention_temperature: float = 1.0
   activation_function_name: str = 'swish'
-  use_post_layer_norm: bool = False
+  use_post_layer_norm: bool = True
 
 
 class LayerNorm(nn.Module):

diff --git a/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/workload.py b/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_jax/workload.py
@@ -377,7 +377,7 @@ class LibriSpeechConformerLayerNormWorkload(LibriSpeechConformerWorkload):
 
   property
   def use_post_layer_norm(self) -> bool:
-    return True
+    return False
 
 
 class LibriSpeechConformerGeluWorkload(LibriSpeechConformerWorkload):

diff --git a/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/models.py b/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/models.py
@@ -44,7 +44,7 @@ class ConformerConfig:
   use_specaug: bool = True
   attention_temperature: float = 1.0
   activation_function_name: str = 'swish'
-  use_post_layer_norm: bool = False
+  use_post_layer_norm: bool = True
 
 
 def initialize(m):

diff --git a/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/workload.py b/algorithmic_efficiency/workloads/librispeech_conformer/librispeech_pytorch/workload.py
@@ -79,7 +79,7 @@ def init_model_fn(
             use_post_layer_norm=self.use_post_layer_norm,
             activation_function_name=activation_function_name))
     self.ctc_loss = torch.nn.CTCLoss(blank=0, reduction='none')
-    model.initialize(model)
+    models.initialize(model)
     self._param_shapes = param_utils.pytorch_param_shapes(model)
     self._param_types = param_utils.pytorch_param_types(self._param_shapes)
     model.to(DEVICE)
@@ -343,7 +343,7 @@ class LibriSpeechConformerLayerNormWorkload(LibriSpeechConformerWorkload):
 
   property
   def use_post_layer_norm(self) -> bool:
-    return True
+    return False
 
 
 class LibriSpeechConformerGeluWorkload(LibriSpeechConformerWorkload):