FutureWarning using torch.load with torch>2.4, torch.serialization.add_safe_globals does not work for torch_geometric.data.Data

🐛 Describe the bug

Hello,

I wanted to report on a warning related to the latest pytorch versions, which may become an issue moving forward.

Since I've moved to pytorch version >2.4, doing torch.save and torch.load of a torch_geometric.data.Data object results in the following warning:

FutureWarning: You are using torch.load with weights_only=False (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for weights_only will be flipped to True. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via torch.serialization.add_safe_globals. We recommend you start setting weights_only=True for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.

This can be reproduced for instance by running

import torch
from torch_geometric.data import Data

data = Data(x=torch.randn(10))
torch.save(data, 'data.pt')
data = torch.load('data.pt')

However, if I do as suggested, that is using torch.serialization.add_safe_globals to whitelist Data and adding the weights_only option in the torch.load call, i.e.

import torch
from torch_geometric.data import Data

torch.serialization.add_safe_globals([Data])

data = Data(x=torch.randn(10))
torch.save(data, 'data.pt')
data = torch.load('data.pt', weights_only=True)

I get the following error

UnpicklingError: Weights only load failed. This file can still be loaded, to do so you have two options 
	(1) 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 torch_geometric.data.data.DataEdgeAttr was not an allowed global by default. Please use `torch.serialization.add_safe_globals([DataEdgeAttr])` to allowlist this global if you trust this class/function.

I may be mistaken, but I think it is not intended that I also add DataEdgeAttr to the serialization whitelist.

This is clearly not a greatly concerning bug right now, as torch.load still works, but a fix may become necessary in the future.

Versions

Collecting environment information...
PyTorch version: 2.4.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Professionnel (10.0.26100 64-bit)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.4 | packaged by Anaconda, Inc. | (main, Jul  5 2023, 13:47:18) [MSC v.1916 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.26100-SP0
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA RTX A1000 Laptop GPU
Nvidia driver version: 556.12
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Name: 12th Gen Intel(R) Core(TM) i7-12700H
Manufacturer: GenuineIntel
Family: 198
Architecture: 9
ProcessorType: 3
DeviceID: CPU0
CurrentClockSpeed: 2300
MaxClockSpeed: 2300
L2CacheSize: 11776
L2CacheSpeed: None
Revision: None

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.0
[pip3] onnx==1.14.1
[pip3] onnxruntime-gpu==1.16.0
[pip3] optree==0.11.0
[pip3] pytorch-ignite==0.4.12
[pip3] torch==2.4.1+cu118
[pip3] torch_cluster==1.6.3+pt24cu118
[pip3] torch-geometric==2.6.1
[pip3] torch_scatter==2.1.2+pt24cu118
[pip3] torch_sparse==0.6.18+pt24cu118
[pip3] torch_spline_conv==1.2.2+pt24cu118
[pip3] torch-tb-profiler==0.4.3
[pip3] torchaudio==2.4.1+cu118
[pip3] torchmetrics==1.2.0
[pip3] torchvision==0.19.1+cu118
[conda] blas                      1.0                         mkl
[conda] mkl                       2023.1.0         h6b88ed4_46357
[conda] mkl-service               2.4.0           py311h2bbff1b_1
[conda] mkl_fft                   1.3.8           py311h2bbff1b_0
[conda] mkl_random                1.2.4           py311h59b6b97_0
[conda] numpy                     1.26.0          py311hdab7c0b_0
[conda] numpy-base                1.26.0          py311hd01c5d8_0
[conda] optree                    0.11.0                   pypi_0    pypi
[conda] pytorch-ignite            0.4.12                   pypi_0    pypi
[conda] torch                     2.4.1+cu118              pypi_0    pypi
[conda] torch-cluster             1.6.3+pt24cu118          pypi_0    pypi
[conda] torch-geometric           2.6.1                    pypi_0    pypi
[conda] torch-scatter             2.1.2+pt24cu118          pypi_0    pypi
[conda] torch-sparse              0.6.18+pt24cu118          pypi_0    pypi
[conda] torch-spline-conv         1.2.2+pt24cu118          pypi_0    pypi
[conda] torch-tb-profiler         0.4.3                    pypi_0    pypi
[conda] torchaudio                2.4.1+cu118              pypi_0    pypi
[conda] torchmetrics              1.2.0                    pypi_0    pypi
[conda] torchvision               0.19.1+cu118             pypi_0    pypi

Start with the provided torch.save and torch.load reproduction using torch_geometric.data.Data, then inspect PyTorch's weights_only loading and safe-global handling for Data and DataEdgeAttr. Done means the reproduced Data object can be loaded with weights_only=True without the reported unsupported-global error, while preserving the serialization behavior described in the issue.