From 79fb2414db117359df541ac7a519c30a91b1b17d Mon Sep 17 00:00:00 2001 From: nileshpatil6 Date: Sat, 4 Jul 2026 23:29:16 +0530 Subject: [PATCH] fix: restore weights_only=False when reloading best checkpoint PR #15314 dropped the explicit weights_only=False from the torch.load call in NeMoModelCheckpoint.on_save_checkpoint (used when save_best_model=True). Since torch 2.6 defaults weights_only to True, this call now fails whenever the checkpoint's hparams contain an OmegaConf DictConfig, which is the normal case for any Hydra-configured model. The checkpoint being loaded here was just written to disk by this same training run a few lines above, so it is a trusted file and loading it with weights_only=False is safe. Fixes #15709 Signed-off-by: nileshpatil6 --- nemo/utils/callbacks/nemo_model_checkpoint.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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