Apply isort and black reformatting

Signed-off-by: guyueh1 <[email protected]>
guyueh1 · Jan 23, 2025 · b18ac96 · b18ac96
1 parent 441036c
commit b18ac96
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/nemo/lightning/pytorch/callbacks/megatron_comm_overlap.py b/nemo/lightning/pytorch/callbacks/megatron_comm_overlap.py
@@ -208,10 +208,13 @@ def _override_user_cfgs(self, comm_overlap_cfg):
                     setattr(comm_overlap_cfg, field.name, user_value)
 
         return comm_overlap_cfg
-    
+
     def _set_num_cuda_device_max_connections(self):
+        import os
+
+        import torch
+
         from nemo.utils import AppState
-        import os, torch
 
         app_state = AppState()
         tp_size = app_state.tensor_model_parallel_size
@@ -259,7 +262,7 @@ def setup(self, trainer: pl.Trainer, pl_module: pl.LightningModule, stage: str)
             self._apply_cfgs(comm_overlap_cfg, trainer.strategy.ddp_config)
             if hasattr(trainer.model, '__io__'):
                 self._apply_cfgs(comm_overlap_cfg, trainer.model.__io__.optim.config)
-        
+
         # setup cuda device max connections
         self._set_num_cuda_device_max_connections()