[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
NVIDIA · Apr 4, 2023 · 988532b · 988532b
1 parent 895a2de
commit 988532b
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py b/nemo/collections/nlp/models/language_modeling/megatron_gpt_model.py
@@ -335,11 +335,10 @@ def training_step(self, dataloader_iter, batch_idx):
         if self.initialize_ub:
             input_shape = [
                 self.cfg.get('encoder_seq_length') * self.cfg.get('micro_batch_size'),
-                self.cfg.get('hidden_size')
+                self.cfg.get('hidden_size'),
             ]
             te_module.pre_init_ub(
-                shape=input_shape,
-                dtype=torch.uint8 if self.cfg.get('fp8') else torch.bfloat16,
+                shape=input_shape, dtype=torch.uint8 if self.cfg.get('fp8') else torch.bfloat16,
             )
             self.initialize_ub = False
 

diff --git a/nemo/collections/nlp/modules/common/megatron/transformer.py b/nemo/collections/nlp/modules/common/megatron/transformer.py
@@ -779,7 +779,7 @@ def __init__(
         layer_type: str = "encoder",
         drop_path_rate: float = 0,
         use_emha: bool = False,
-        ub_tp_comm_overlap: bool =False,
+        ub_tp_comm_overlap: bool = False,
         autocast_dtype: Any = 16,
         zero_centered_gamma: bool = False,
     ) -> None: