Remove mask if use fusion mask (NVIDIA#9723)

* Remove mask if use fusion mask Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> * Apply isort and black reformatting Signed-off-by: hsiehjackson <hsiehjackson@users.noreply.github.com> --------- Signed-off-by: Cheng-Ping Hsieh <chsieh@nvidia.com> Signed-off-by: hsiehjackson <hsiehjackson@users.noreply.github.com> Co-authored-by: hsiehjackson <hsiehjackson@users.noreply.github.com>
nikitaved · Jul 16, 2024 · e28fc39 · e28fc39
1 parent 2fed71c
commit e28fc39
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/nemo/collections/nlp/modules/common/text_generation_utils.py b/nemo/collections/nlp/modules/common/text_generation_utils.py
@@ -794,6 +794,9 @@ def generate(
     if random_seed is not None:
         seed_everything(random_seed)
 
+    if hasattr(model, 'get_attention_mask_from_fusion') and model.get_attention_mask_from_fusion:
+        compute_attention_mask = False
+
     output = synced_generate(
         model,
         inference_strategy,