Merge branch 'main' into jiemingz/ckpt_mem_fixes

JimmyZhang12 · Mar 26, 2024 · c4545fe · c4545fe
2 parents 3d8dd95 + 439c14b
commit c4545fe
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/nemo/core/classes/modelPT.py b/nemo/core/classes/modelPT.py
@@ -1732,7 +1732,7 @@ def on_train_batch_start(self, batch: Any, batch_idx: int, unused: int = 0) -> O
         if self.device.type == 'cuda':
             if hasattr(self, '_nsys_profile_enabled'):
                 if self._nsys_profile_enabled and not self._profile_complete:
-                    if batch_idx == self._nsys_profile_start_step and get_rank() in self._nsys_profile_ranks:
+                    if batch_idx >= self._nsys_profile_start_step and get_rank() in self._nsys_profile_ranks:
                         logging.info("====== Start nsys profiling ======")
                         torch.cuda.cudart().cudaProfilerStart()
                         if self._nsys_profile_gen_shape:
@@ -1769,7 +1769,7 @@ def on_train_batch_end(self, outputs, batch: Any, batch_idx: int, unused: int =
         if self.device.type == 'cuda':
             if hasattr(self, '_nsys_profile_enabled'):
                 if self._nsys_profile_enabled and not self._profile_complete:
-                    if batch_idx == self._nsys_profile_end_step and get_rank() in self._nsys_profile_ranks:
+                    if batch_idx >= self._nsys_profile_end_step and get_rank() in self._nsys_profile_ranks:
                         logging.info("====== End nsys profiling ======")
                         torch.cuda.cudart().cudaProfilerStop()
                         self._profile_complete = True