Fix OOM

Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com>
NVIDIA · Aug 18, 2023 · ae3f7d2 · ae3f7d2
1 parent ef0e981
commit ae3f7d2
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/nemo/collections/nlp/modules/common/megatron/megatron_transformer_encoder.py b/nemo/collections/nlp/modules/common/megatron/megatron_transformer_encoder.py
@@ -170,10 +170,10 @@ def forward(
         # convert to Megatron mask
         if self.use_flash_attention:
             enc_attn_mask_3d = enc_attn_mask < 0.5
-
-        enc_attn_mask_3d = attn_mask_postprocess(build_attention_mask_3d(
-            source_mask=enc_attn_mask, target_mask=enc_attn_mask, attn_mask_type=self.model_attn_mask_type,
-        ))
+        else:
+            enc_attn_mask_3d = attn_mask_postprocess(build_attention_mask_3d(
+                source_mask=enc_attn_mask, target_mask=enc_attn_mask, attn_mask_type=self.model_attn_mask_type,
+            ))
 
         # transformer encoder
         enc_output = self.model(