fix PygLinkPropPredDataset for torch.load weights_only in latest pytorch

[yifanguan]

Starting in PyTorch 2.6, the default value of weights_only argument of torch.load is changed from False to True, causing loading ogb link predication pyg datasets to fail. This pull request fix this issue by explicitly providing weights_only=True when loading link predication pyg-version datasets.

Below is more details:
Traceback (most recent call last):
File "/home/yguan10/GraphLearning/link_predication/link_predication.py", line 263, in
main()
File "/home/yguan10/GraphLearning/link_predication/link_predication.py", line 196, in main
split_edge = dataset.get_edge_split()
^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/yguan10/GraphLearning/venv/lib/python3.12/site-packages/ogb/linkproppred/dataset_pyg.py", line 77, in get_edge_split
train = replace_numpy_with_torchtensor(torch.load(osp.join(path, 'train.pt')))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/yguan10/GraphLearning/venv/lib/python3.12/site-packages/torch/serialization.py", line 1548, in load
raise pickle.UnpicklingError(_get_wo_message(str(e))) from None
_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.
(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 will likely succeed, but it can result in arbitrary code execution. Do it only if you got the file from a trusted source.
(2) Alternatively, to load with weights_only=True please check the recommended steps in the following error message.
WeightsUnpickler error: Unsupported global: GLOBAL numpy.core.multiarray._reconstruct was not an allowed global by default. Please use torch.serialization.add_safe_globals([numpy.core.multiarray._reconstruct]) or the torch.serialization.safe_globals([numpy.core.multiarray._reconstruct]) context manager to allowlist this global if you trust this class/function.

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.

[mathis-saunier]

Hi @yifanguan, I stumbled across the same issue.

I noticed the same issue and was looking at your PR.

About this line:
return torch.load(os.path.join(path, 'split_dict.pt', weights_only=False))

It looks like weights_only=False ended up inside os.path.join(), which will cause an error since this argument belongs to torch.load(). Did you intend something like:
return torch.load(os.path.join(path, 'split_dict.pt'), weights_only=False)

Maybe I am wrong. Did you run some basic tests with your changes ?

Regarding the CI failure: GitHub’s pre-commit.ci fails because the repository doesn’t include a .pre-commit-config.yaml file. Adding one at the root of the project should fix the check. You can find the setup instructions here:
https://pre-commit.com/

If you have any questions, feel free to ask.
Sincerely,