diff --git a/nemo/core/classes/modelPT.py b/nemo/core/classes/modelPT.py index 2c337a6a6f48..754d18eb8267 100644 --- a/nemo/core/classes/modelPT.py +++ b/nemo/core/classes/modelPT.py @@ -1760,7 +1760,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 not self._profile_complete: + 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: logging.info("====== End nsys profiling ======") torch.cuda.cudart().cudaProfilerStop()