Bug: UnpicklingError When Loading the Checkpoint on Rank 0

[wyin1632]

I got the following error during a model sweep.

[rank0]:   File "/home/ubuntu/miniconda3/envs/torch_tabular/lib/python3.10/site-packages/pytorch_tabular/utils/python_utils.py", line 85, in pl_load
[rank0]:     return torch.load(f, map_location=map_location)
[rank0]:   File "/home/ubuntu/miniconda3/envs/torch_tabular/lib/python3.10/site-packages/torch/serialization.py", line 1470, in load
[rank0]:     raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
[rank0]: _pickle.UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options, do those steps only if you trust the source of the checkpoint. 
[rank0]:        (1) In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` w
ill likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
[rank0]:        (2) Alternatively, to load with `weights_only=True` please check the recommended steps in the following error message.
[rank0]:        WeightsUnpickler error: Unsupported global: GLOBAL omegaconf.dictconfig.DictConfig was not an allowed global by default. Please use `torch.serialization.add_safe_globals([DictConfi
g])` or the `torch.serialization.safe_globals([DictConfig])` context manager to allowlist this global if you trust this class/function.

[rank0]: Check the documentation of torch.load to learn more about types accepted by default with weights_only https://pytorch.org/docs/stable/generated/torch.load.html.

Chaging to torch.load(f, map_location=map_location, weights_only=False) solved the error. Should we set weights_only=False in the code by default?

[dosubot]

Yes, setting weights_only=False is recommended to solve the UnpicklingError when loading the checkpoint on Rank 0. Additionally, you can address this issue by adding the following code to allow specific globals:

from omegaconf import DictConfig
torch.serialization.safe_globals([DictConfig])

This should be done alongside setting weights_only=False when using torch.load(). Here is the relevant code:

torch.load(path_to_checkpoint, map_location=map_location, weights_only=False)

This solution is effective across different environments, including Docker containers with Python 3.12.7 on Linux Debian. However, it should only be done if you trust the source of the checkpoint, as it can result in arbitrary code execution ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[manujosephv]

Pytorch 2.6 updated the torch.load function with "weights_only = True" as default. PR #543 is fixing this.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.