diff --git a/nemo/utils/callbacks/nemo_model_checkpoint.py b/nemo/utils/callbacks/nemo_model_checkpoint.py index 1285bece3cc4..8752b6f30d84 100644 --- a/nemo/utils/callbacks/nemo_model_checkpoint.py +++ b/nemo/utils/callbacks/nemo_model_checkpoint.py @@ -237,7 +237,11 @@ def on_save_checkpoint(self, trainer, pl_module, checkpoint): self.previous_best_path = self.best_model_path old_state_dict = deepcopy(pl_module.state_dict()) - checkpoint = torch.load(maybe_injected_best_model_path, map_location='cpu') + # This checkpoint was just written to disk by this same training run (see + # `_save_checkpoint` below), so it is trusted. It also contains the model's + # hparams (e.g. OmegaConf DictConfig), which `weights_only=True` (the default + # since torch 2.6) cannot unpickle, so it must be loaded with weights_only=False. + checkpoint = torch.load(maybe_injected_best_model_path, map_location='cpu', weights_only=False) if 'state_dict' in checkpoint: checkpoint = checkpoint['state_dict'] # get a new instanace of the model