[Python] torch cannot add NamedTuple class to safe_globals

I define some named tuple like this:

class checkpoint_t(NamedTuple):
epoch: int
model_state_dict: Dict[str, Any]
optimizer_state_dict: Dict[str, Any]
model_name: str | None = None



However I after save I cannot load this namedtuple via

import torch
from train import checkpoint_t
with torch.serialization.safe_globals([checkpoint_t]):
print("safe globals: ", torch.serialization.get_safe_globals())
checkpoint: checkpoint_t = torch.load(parsed_args.checkpoint, weights_only=True)


it's still saying:

WeightsUnpickler error: Unsupported global: GLOBAL main.checkpoint_t was not an allowed global by default. Please use torch.serialization.add_safe_globals([checkpoint_t]) or the torch.serialization.safe_globals([checkpoint_t]) context manager to allowlist this global if you trust this class/function.
​

any idea why? and how to fix this?

Continue reading...