From 9dbedebf3da4c4d9166631038d400391cf36f069 Mon Sep 17 00:00:00 2001 From: alex-crlhmmr Date: Fri, 10 Jul 2026 05:14:55 -0700 Subject: [PATCH] Fix pretrained checkpoint loading under torch < 2.4 load_checkpoint_state_dict guards add_safe_globals with hasattr, but that API does not exist before torch 2.4, so on older torch the guard is a no-op and torch.load(weights_only=True) rejects the argparse.Namespace that Lightning stores under hyper_parameters. Loading a pretraining checkpoint into the finetuning path then fails with an UnpicklingError. Fall back to weights_only=False when add_safe_globals is unavailable. The checkpoint is a local file produced by pretrain.py, so a full load is safe. --- models/checkpoint.py | 22 ++++++++++++++++------ 1 file changed, 16 insertions(+), 6 deletions(-) diff --git a/models/checkpoint.py b/models/checkpoint.py index a099ed1..75647a6 100644 --- a/models/checkpoint.py +++ b/models/checkpoint.py @@ -38,12 +38,22 @@ def load_checkpoint_state_dict( torch.serialization.add_safe_globals( [argparse.Namespace, pathlib.PosixPath, pathlib.WindowsPath] ) - - checkpoint = torch.load( - checkpoint_path, - map_location="cpu", - weights_only=True, - ) + checkpoint = torch.load( + checkpoint_path, + map_location="cpu", + weights_only=True, + ) + else: + # torch < 2.4 has no add_safe_globals, so the weights_only unpickler + # cannot be extended to allow the argparse.Namespace that Lightning + # stores under hyper_parameters. A pretraining checkpoint therefore + # fails to load with weights_only=True. Fall back to a full load; the + # checkpoint is a local, trusted file produced by pretrain.py. + checkpoint = torch.load( + checkpoint_path, + map_location="cpu", + weights_only=False, + ) if not isinstance(checkpoint, Mapping): raise TypeError("Checkpoint must contain a mapping")