From 45b7e0779fa22db46398f4de9cb7f5a51af1296a Mon Sep 17 00:00:00 2001 From: mikolajblaz Date: Wed, 15 May 2024 13:57:43 +0200 Subject: [PATCH] Fix incorrect checkpoint removal logic (#9192) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Fix incorrect if logic Signed-off-by: Mikołaj Błaż * Apply isort and black reformatting Signed-off-by: mikolajblaz --------- Signed-off-by: Mikołaj Błaż Signed-off-by: mikolajblaz --- nemo/collections/nlp/parts/nlp_overrides.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/nemo/collections/nlp/parts/nlp_overrides.py b/nemo/collections/nlp/parts/nlp_overrides.py index 1c68ebff8121..4e262912ccc3 100644 --- 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: