clean up

Signed-off-by: Sangkug Lym <[email protected]>
NVIDIA · Dec 2, 2023 · 2d64e8e · 2d64e8e
1 parent 29826c9
commit 2d64e8e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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()