Fix layer_norm_1p Apex call to match Apex API

Signed-off-by: Shriya Palsamudram <spalsamudram@nvidia.com>
NVIDIA · Jan 9, 2024 · d210a75 · d210a75
1 parent 76a712a
commit d210a75
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/modules/common/megatron/layer_norm_1p.py b/nemo/collections/nlp/modules/common/megatron/layer_norm_1p.py
@@ -40,7 +40,7 @@ def reset_parameters(self):
             torch.nn.init.zeros_(self.bias)
 
         def forward(self, x):
-            return _fast_layer_norm(x, self.weight + 1, self.bias, self.epsilon)
+            return _fast_layer_norm(x, self.weight + 1, self.bias, self.epsilon, self.memory_efficient)
 
 
 else: