Fix incorrect checkpoint removal logic (#9192)

* Fix incorrect if logic Signed-off-by: Mikołaj Błaż <[email protected]> * Apply isort and black reformatting Signed-off-by: mikolajblaz <[email protected]> --------- Signed-off-by: Mikołaj Błaż <[email protected]> Signed-off-by: mikolajblaz <[email protected]>
NVIDIA · May 15, 2024 · 45b7e07 · 45b7e07
1 parent bb26e98
commit 45b7e07
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py
@@ -452,8 +452,9 @@ def load_checkpoint(self, checkpoint_path: Union[str, Path]) -> Dict[str, Any]:
 
     def remove_checkpoint(self, filepath: Union[str, Path]) -> None:
         # check if filepath is a distributed checkpoint
-        if self.use_distributed_checkpointing and self.is_global_zero:
-            self.checkpoint_io.remove_checkpoint(ckpt_to_dir(filepath))
+        if self.use_distributed_checkpointing:
+            if self.is_global_zero:
+                self.checkpoint_io.remove_checkpoint(ckpt_to_dir(filepath))
 
         # legacy checkpoint logic, does not use megatron core
         else: