best ckpt fix (NVIDIA#7564) (NVIDIA#7588)

Signed-off-by: dimapihtar <[email protected]> Co-authored-by: Dmytro Pykhtar <[email protected]> Signed-off-by: Sasha Meister <[email protected]>
ssh-meister · Oct 10, 2023 · 71f327f · 71f327f
1 parent 2f6fa29
commit 71f327f
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/nemo/utils/callbacks/nemo_model_checkpoint.py b/nemo/utils/callbacks/nemo_model_checkpoint.py
@@ -209,6 +209,8 @@ def on_train_end(self, trainer, pl_module):
                     "were found. Saving latest model instead."
                 )
             else:
+                if os.path.isdir(self.best_model_path.split('.ckpt')[0]):
+                    self.best_model_path = self.best_model_path.split('.ckpt')[0]
                 self.best_model_path = trainer.strategy.broadcast(self.best_model_path)
                 trainer._checkpoint_connector.restore(self.best_model_path)