TE bug fix (NVIDIA#7027)

Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
borisfom · Jul 14, 2023 · 49b1aea · 49b1aea
1 parent fac7103
commit 49b1aea
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -525,7 +525,7 @@ def initialize_ub_func(self):
                     ub_cfgs = yaml.safe_load(ub_cfg_file)
             except (ImportError, TypeError):
                 logging.error(f"Fail to read ub_tp_comm_overlap config file: {ub_cfg_file_name}.")
-        te_module.initialize_ub(
+        te_module.base.initialize_ub(
             shape=input_shape,
             tp_size=self.cfg.get('tensor_model_parallel_size'),
             use_fp8=self.cfg.get('fp8'),