pyg-nightly 2.6.0.dev20240318__py3-none-any.whl → 2.7.0.dev20250115__py3-none-any.whl

{pyg_nightly-2.6.0.dev20240318.dist-info → pyg_nightly-2.7.0.dev20250115.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: flit 3.9.0
+Generator: flit 3.10.1
 Root-Is-Purelib: true
 Tag: py3-none-any

torch_geometric/__init__.py

@@ -1,7 +1,15 @@
+from collections import defaultdict
+
+import torch
+import torch_geometric.typing
+
 from ._compile import compile, is_compiling
+from ._onnx import is_in_onnx_export
+from .index import Index
 from .edge_index import EdgeIndex
 from .seed import seed_everything
 from .home import get_home_dir, set_home_dir
+from .device import is_mps_available, is_xpu_available, device
 from .isinstance import is_torch_instance
 from .debug import is_debug_enabled, debug, set_debug
 
@@ -22,15 +30,20 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.6.0.dev20240318'
+__version__ = '2.7.0.dev20250115'
 
 __all__ = [
+    'Index',
     'EdgeIndex',
     'seed_everything',
     'get_home_dir',
     'set_home_dir',
     'compile',
     'is_compiling',
+    'is_in_onnx_export',
+    'is_mps_available',
+    'is_xpu_available',
+    'device',
     'is_torch_instance',
     'is_debug_enabled',
     'debug',
@@ -41,3 +54,17 @@ __all__ = [
     'torch_geometric',
     '__version__',
 ]
+
+# Serialization ###############################################################
+
+if torch_geometric.typing.WITH_PT24:
+    torch.serialization.add_safe_globals([
+        dict,
+        list,
+        defaultdict,
+        Index,
+        torch_geometric.index.CatMetadata,
+        EdgeIndex,
+        torch_geometric.edge_index.SortOrder,
+        torch_geometric.edge_index.CatMetadata,
+    ])

torch_geometric/_compile.py

@@ -10,6 +10,8 @@ def is_compiling() -> bool:
     r"""Returns :obj:`True` in case :pytorch:`PyTorch` is compiling via
     :meth:`torch.compile`.
     """
+    if torch_geometric.typing.WITH_PT23:
+        return torch.compiler.is_compiling()
     if torch_geometric.typing.WITH_PT21:
         return torch._dynamo.is_compiling()
     return False  # pragma: no cover
@@ -25,10 +27,15 @@ def compile(
     This function has the same signature as :meth:`torch.compile` (see
     `here <https://pytorch.org/docs/stable/generated/torch.compile.html>`__).
 
+    Args:
+        model: The model to compile.
+        *args: Additional arguments of :meth:`torch.compile`.
+        **kwargs: Additional keyword arguments of :meth:`torch.compile`.
+
     .. note::
         :meth:`torch_geometric.compile` is deprecated in favor of
         :meth:`torch.compile`.
     """
     warnings.warn("'torch_geometric.compile' is deprecated in favor of "
                   "'torch.compile'")
-    return torch.compile(model, *args, **kwargs)
+    return torch.compile(model, *args, **kwargs)  # type: ignore

torch_geometric/_onnx.py

@@ -0,0 +1,14 @@
+import torch
+
+from torch_geometric import is_compiling
+
+
+def is_in_onnx_export() -> bool:
+    r"""Returns :obj:`True` in case :pytorch:`PyTorch` is exporting to ONNX via
+    :meth:`torch.onnx.export`.
+    """
+    if is_compiling():
+        return False
+    if torch.jit.is_scripting():
+        return False
+    return torch.onnx.is_in_onnx_export()

torch_geometric/config_mixin.py

@@ -0,0 +1,113 @@
+import inspect
+from dataclasses import fields, is_dataclass
+from importlib import import_module
+from typing import Any, Dict
+
+from torch_geometric.config_store import (
+    class_from_dataclass,
+    dataclass_from_class,
+)
+from torch_geometric.isinstance import is_torch_instance
+
+
+class ConfigMixin:
+    r"""Enables a class to serialize/deserialize itself to a dataclass."""
+    def config(self) -> Any:
+        r"""Creates a serializable configuration of the class."""
+        data_cls = dataclass_from_class(self.__class__)
+        if data_cls is None:
+            raise ValueError(f"Could not find the configuration class that "
+                             f"belongs to '{self.__class__.__name__}'. Please "
+                             f"register it in the configuration store.")
+
+        kwargs: Dict[str, Any] = {}
+        for field in fields(data_cls):
+            if not hasattr(self, field.name):
+                continue
+            kwargs[field.name] = _recursive_config(getattr(self, field.name))
+        return data_cls(**kwargs)
+
+    @classmethod
+    def from_config(cls, cfg: Any, *args: Any, **kwargs: Any) -> Any:
+        r"""Instantiates the class from a serializable configuration."""
+        if getattr(cfg, '_target_', None):
+            cls = _locate_cls(cfg._target_)
+        elif isinstance(cfg, dict) and '_target_' in cfg:
+            cls = _locate_cls(cfg['_target_'])
+
+        data_cls = cfg.__class__
+        if not is_dataclass(data_cls):
+            data_cls = dataclass_from_class(cls)
+            if data_cls is None:
+                raise ValueError(f"Could not find the configuration class "
+                                 f"that belongs to '{cls.__name__}'. Please "
+                                 f"register it in the configuration store.")
+
+        field_names = {field.name for field in fields(data_cls)}
+        if isinstance(cfg, dict):
+            _kwargs = {k: v for k, v in cfg.items() if k in field_names}
+            cfg = data_cls(**_kwargs)
+        assert is_dataclass(cfg)
+
+        if len(args) > 0:  # Convert `*args` to `**kwargs`:
+            param_names = list(inspect.signature(cls).parameters.keys())
+            if 'args' in param_names:
+                param_names.remove('args')
+            if 'kwargs' in param_names:
+                param_names.remove('kwargs')
+
+            for name, arg in zip(param_names, args):
+                kwargs[name] = arg
+
+        for key in field_names:
+            if key not in kwargs and key != '_target_':
+                kwargs[key] = _recursive_from_config(getattr(cfg, key))
+
+        return cls(**kwargs)
+
+
+def _recursive_config(value: Any) -> Any:
+    if isinstance(value, ConfigMixin):
+        return value.config()
+    if is_torch_instance(value, ConfigMixin):
+        return value.config()
+    if isinstance(value, (tuple, list)):
+        return [_recursive_config(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _recursive_config(v) for k, v in value.items()}
+    return value
+
+
+def _recursive_from_config(value: Any) -> Any:
+    cls: Any = None
+    if is_dataclass(value):
+        if getattr(value, '_target_', None):
+            try:
+                cls = _locate_cls(value._target_)  # type: ignore
+            except ImportError:
+                pass  # Keep the dataclass as it is.
+        else:
+            cls = class_from_dataclass(value.__class__)
+    elif isinstance(value, dict) and '_target_' in value:
+        cls = _locate_cls(value['_target_'])
+
+    if cls is not None and issubclass(cls, ConfigMixin):
+        return cls.from_config(value)
+    if isinstance(value, (tuple, list)):
+        return [_recursive_from_config(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _recursive_from_config(v) for k, v in value.items()}
+    return value
+
+
+def _locate_cls(qualname: str) -> Any:
+    parts = qualname.split('.')
+
+    if len(parts) <= 1:
+        raise ValueError(f"Qualified name is missing a dot (got '{qualname}')")
+
+    if any([len(part) == 0 for part in parts]):
+        raise ValueError(f"Relative imports not supported (got '{qualname}')")
+
+    module_name, cls_name = '.'.join(parts[:-1]), parts[-1]
+    return getattr(import_module(module_name), cls_name)

torch_geometric/config_store.py

@@ -76,7 +76,7 @@ else:
 
         def __call__(cls, *args: Any, **kwargs: Any) -> Any:
             if cls not in cls._instances:
-                instance = super(Singleton, cls).__call__(*args, **kwargs)
+                instance = super().__call__(*args, **kwargs)
                 cls._instances[cls] = instance
                 return instance
             return cls._instances[cls]
@@ -162,12 +162,19 @@ def map_annotation(
     annotation: Any,
     mapping: Optional[Dict[Any, Any]] = None,
 ) -> Any:
-
     origin = getattr(annotation, '__origin__', None)
-    args = getattr(annotation, '__args__', [])
-    if origin == Union or origin == list or origin == dict:
-        annotation = copy.copy(annotation)
-        annotation.__args__ = tuple(map_annotation(a, mapping) for a in args)
+    args: Tuple[Any, ...] = getattr(annotation, '__args__', tuple())
+    if origin in {Union, list, dict, tuple}:
+        assert origin is not None
+        args = tuple(map_annotation(a, mapping) for a in args)
+        if type(annotation).__name__ == 'GenericAlias':
+            # If annotated with `list[...]` or `dict[...]` (>= Python 3.10):
+            annotation = origin[args]
+        else:
+            # If annotated with `typing.List[...]` or `typing.Dict[...]`:
+            annotation = copy.copy(annotation)
+            annotation.__args__ = args
+
         return annotation
 
     if mapping is not None and annotation in mapping:
@@ -231,7 +238,7 @@ def to_dataclass(
     if strict:  # Check that keys in map_args or exclude_args are present.
         keys = set() if map_args is None else set(map_args.keys())
         if exclude_args is not None:
-            keys |= set([arg for arg in exclude_args if isinstance(arg, str)])
+            keys |= {arg for arg in exclude_args if isinstance(arg, str)}
         diff = keys - set(params.keys())
         if len(diff) > 0:
             raise ValueError(f"Expected input argument(s) {diff} in "
@@ -406,13 +413,13 @@ def fill_config_store() -> None:
 
     # Register `torch_geometric.transforms` ###################################
     transforms = torch_geometric.transforms
-    for cls_name in set(transforms.__all__) - set([
+    for cls_name in set(transforms.__all__) - {
             'BaseTransform',
             'Compose',
             'ComposeFilters',
             'LinearTransformation',
             'AddMetaPaths',  # TODO
-    ]):
+    }:
         cls = to_dataclass(getattr(transforms, cls_name), base_cls=Transform)
         # We use an explicit additional nesting level inside each config to
         # allow for applying multiple transformations.
@@ -426,7 +433,7 @@ def fill_config_store() -> None:
         'pre_transform': (Dict[str, Transform], field(default_factory=dict)),
     }
 
-    for cls_name in set(datasets.__all__) - set([]):
+    for cls_name in set(datasets.__all__) - set():
         cls = to_dataclass(getattr(datasets, cls_name), base_cls=Dataset,
                            map_args=map_dataset_args,
                            exclude_args=['pre_filter'])
@@ -434,32 +441,34 @@ def fill_config_store() -> None:
 
     # Register `torch_geometric.models` #######################################
     models = torch_geometric.nn.models.basic_gnn
-    for cls_name in set(models.__all__) - set([]):
+    for cls_name in set(models.__all__) - set():
         cls = to_dataclass(getattr(models, cls_name), base_cls=Model)
         config_store.store(cls_name, group='model', node=cls)
 
     # Register `torch.optim.Optimizer` ########################################
-    for cls_name in set([
-            key for key, cls in torch.optim.__dict__.items()
+    for cls_name in {
+            key
+            for key, cls in torch.optim.__dict__.items()
             if inspect.isclass(cls) and issubclass(cls, torch.optim.Optimizer)
-    ]) - set([
+    } - {
             'Optimizer',
-    ]):
+    }:
         cls = to_dataclass(getattr(torch.optim, cls_name), base_cls=Optimizer,
                            exclude_args=['params'])
         config_store.store(cls_name, group='optimizer', node=cls)
 
     # Register `torch.optim.lr_scheduler` #####################################
-    for cls_name in set([
-            key for key, cls in torch.optim.lr_scheduler.__dict__.items()
+    for cls_name in {
+            key
+            for key, cls in torch.optim.lr_scheduler.__dict__.items()
             if inspect.isclass(cls)
-    ]) - set([
+    } - {
             'Optimizer',
             '_LRScheduler',
             'Counter',
             'SequentialLR',
             'ChainedScheduler',
-    ]):
+    }:
         cls = to_dataclass(getattr(torch.optim.lr_scheduler, cls_name),
                            base_cls=LRScheduler, exclude_args=['optimizer'])
         config_store.store(cls_name, group='lr_scheduler', node=cls)

torch_geometric/data/__init__.py

@@ -1,7 +1,10 @@
 # flake8: noqa
 
+import torch
+import torch_geometric.typing
+
 from .feature_store import FeatureStore, TensorAttr
-from .graph_store import GraphStore, EdgeAttr
+from .graph_store import GraphStore, EdgeAttr, EdgeLayout
 from .data import Data
 from .hetero_data import HeteroData
 from .batch import Batch
@@ -13,6 +16,7 @@ from .on_disk_dataset import OnDiskDataset
 from .makedirs import makedirs
 from .download import download_url, download_google_url
 from .extract import extract_tar, extract_zip, extract_bz2, extract_gz
+from .large_graph_indexer import LargeGraphIndexer, TripletLike, get_features_for_triplets, get_features_for_triplets_groups
 
 from torch_geometric.lazy_loader import LazyLoader
 
@@ -24,6 +28,8 @@ data_classes = [
     'Dataset',
     'InMemoryDataset',
     'OnDiskDataset',
+    'LargeGraphIndexer',
+    'TripletLike',
 ]
 
 remote_backend_classes = [
@@ -47,6 +53,8 @@ helper_functions = [
     'extract_zip',
     'extract_bz2',
     'extract_gz',
+    'get_features_for_triplets',
+    "get_features_for_triplets_groups",
 ]
 
 __all__ = data_classes + remote_backend_classes + helper_functions
@@ -68,6 +76,21 @@ from torch_geometric.loader import DataLoader
 from torch_geometric.loader import DataListLoader
 from torch_geometric.loader import DenseDataLoader
 
+# Serialization ###############################################################
+
+if torch_geometric.typing.WITH_PT24:
+    torch.serialization.add_safe_globals([
+        Data,
+        HeteroData,
+        TemporalData,
+        ClusterData,
+        TensorAttr,
+        EdgeAttr,
+        EdgeLayout,
+    ])
+
+# Deprecations ################################################################
+
 NeighborSampler = deprecated(  # type: ignore
     details="use 'loader.NeighborSampler' instead",
     func_name='data.NeighborSampler',

torch_geometric/data/batch.py

@@ -118,8 +118,8 @@ class Batch(metaclass=DynamicInheritance):
         """
         if not hasattr(self, '_slice_dict'):
             raise RuntimeError(
-                ("Cannot reconstruct 'Data' object from 'Batch' because "
-                 "'Batch' was not created via 'Batch.from_data_list()'"))
+                "Cannot reconstruct 'Data' object from 'Batch' because "
+                "'Batch' was not created via 'Batch.from_data_list()'")
 
         data = separate(
             cls=self.__class__.__bases__[-1],

torch_geometric/data/collate.py

@@ -16,7 +16,7 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage, NodeStorage
 from torch_geometric.edge_index import SortOrder
@@ -184,7 +184,8 @@ def _collate(
             return value, slices, incs
 
         out = None
-        if torch.utils.data.get_worker_info() is not None:
+        if (torch.utils.data.get_worker_info() is not None
+                and not isinstance(elem, (Index, EdgeIndex))):
             # Write directly into shared memory to avoid an extra copy:
             numel = sum(value.numel() for value in values)
             if torch_geometric.typing.WITH_PT20:
@@ -203,6 +204,11 @@ def _collate(
 
         value = torch.cat(values, dim=cat_dim or 0, out=out)
 
+        if increment and isinstance(value, Index) and values[0].is_sorted:
+            # Check whether the whole `Index` is sorted:
+            if (value.diff() >= 0).all():
+                value._is_sorted = True
+
         if increment and isinstance(value, EdgeIndex) and values[0].is_sorted:
             # Check whether the whole `EdgeIndex` is sorted by row:
             if values[0].is_sorted_by_row and (value[0].diff() >= 0).all():

torch_geometric/data/data.py

@@ -31,6 +31,7 @@ from torch_geometric.data.storage import (
     NodeStorage,
 )
 from torch_geometric.deprecation import deprecated
+from torch_geometric.index import Index
 from torch_geometric.typing import (
     EdgeTensorType,
     EdgeType,
@@ -290,13 +291,14 @@ class BaseData:
         self,
         start_time: Union[float, int],
         end_time: Union[float, int],
+        attr: str = 'time',
     ) -> Self:
         r"""Returns a snapshot of :obj:`data` to only hold events that occurred
         in period :obj:`[start_time, end_time]`.
         """
         out = copy.copy(self)
         for store in out.stores:
-            store.snapshot(start_time, end_time)
+            store.snapshot(start_time, end_time, attr)
         return out
 
     def up_to(self, end_time: Union[float, int]) -> Self:
@@ -644,7 +646,7 @@ class Data(BaseData, FeatureStore, GraphStore):
         return self
 
     def __cat_dim__(self, key: str, value: Any, *args, **kwargs) -> Any:
-        if is_sparse(value) and 'adj' in key:
+        if is_sparse(value) and ('adj' in key or 'edge_index' in key):
             return (0, 1)
         elif 'index' in key or key == 'face':
             return -1
@@ -653,9 +655,17 @@ class Data(BaseData, FeatureStore, GraphStore):
 
     def __inc__(self, key: str, value: Any, *args, **kwargs) -> Any:
         if 'batch' in key and isinstance(value, Tensor):
+            if isinstance(value, Index):
+                return value.get_dim_size()
             return int(value.max()) + 1
         elif 'index' in key or key == 'face':
-            return self.num_nodes
+            num_nodes = self.num_nodes
+            if num_nodes is None:
+                raise RuntimeError(f"Unable to infer 'num_nodes' from the "
+                                   f"attribute '{key}'. Please explicitly set "
+                                   f"'num_nodes' as an attribute of 'data' to "
+                                   f"prevent this error")
+            return num_nodes
         else:
             return 0
 
@@ -934,16 +944,14 @@ class Data(BaseData, FeatureStore, GraphStore):
         r"""Iterates over all attributes in the data, yielding their attribute
         names and values.
         """
-        for key, value in self._store.items():
-            yield key, value
+        yield from self._store.items()
 
     def __call__(self, *args: str) -> Iterable:
         r"""Iterates over all attributes :obj:`*args` in the data, yielding
         their attribute names and values.
         If :obj:`*args` is not given, will iterate over all attributes.
         """
-        for key, value in self._store.items(*args):
-            yield key, value
+        yield from self._store.items(*args)
 
     @property
     def x(self) -> Optional[Tensor]:
@@ -1163,7 +1171,7 @@ def size_repr(key: Any, value: Any, indent: int = 0) -> str:
                f'[{value.num_rows}, {value.num_cols}])')
     elif isinstance(value, str):
         out = f"'{value}'"
-    elif isinstance(value, Sequence):
+    elif isinstance(value, (Sequence, set)):
         out = str([len(value)])
     elif isinstance(value, Mapping) and len(value) == 0:
         out = '{}'

torch_geometric/data/database.py

@@ -1,16 +1,15 @@
-import pickle
+import io
 import warnings
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from functools import cached_property
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
-from uuid import uuid4
 
 import torch
 from torch import Tensor
 from tqdm import tqdm
 
-from torch_geometric import EdgeIndex
+from torch_geometric import EdgeIndex, Index
 from torch_geometric.edge_index import SortOrder
 from torch_geometric.utils.mixin import CastMixin
 
@@ -19,9 +18,17 @@ from torch_geometric.utils.mixin import CastMixin
 class TensorInfo(CastMixin):
     dtype: torch.dtype
     size: Tuple[int, ...] = (-1, )
+    is_index: bool = False
     is_edge_index: bool = False
 
     def __post_init__(self) -> None:
+        if self.is_index and self.is_edge_index:
+            raise ValueError("Tensor cannot be a 'Index' and 'EdgeIndex' "
+                             "tensor at the same time")
+
+        if self.is_index:
+            self.size = (-1, )
+
         if self.is_edge_index:
             self.size = (2, -1)
 
@@ -33,7 +40,8 @@ def maybe_cast_to_tensor_info(value: Any) -> Union[Any, TensorInfo]:
         return value
     if 'dtype' not in value:
         return value
-    if len(set(value.keys()) | {'dtype', 'size', 'is_edge_index'}) != 3:
+    valid_keys = {'dtype', 'size', 'is_index', 'is_edge_index'}
+    if len(set(value.keys()) | valid_keys) != len(valid_keys):
         return value
     return TensorInfo.cast(value)
 
@@ -107,11 +115,9 @@ class Database(ABC):
         r"""Connects to the database.
         Databases will automatically connect on instantiation.
         """
-        pass
 
     def close(self) -> None:
         r"""Closes the connection to the database."""
-        pass
 
     @abstractmethod
     def insert(self, index: int, data: Any) -> None:
@@ -373,8 +379,9 @@ class SQLiteDatabase(Database):
 
         # We create a temporary ID table to then perform an INNER JOIN.
         # This avoids having a long IN clause and guarantees sorted outputs:
-        join_table_name = f'{self.name}__join__{uuid4().hex}'
-        query = (f'CREATE TABLE {join_table_name} (\n'
+        join_table_name = f'{self.name}__join'
+        # Temporary tables do not lock the database.
+        query = (f'CREATE TEMP TABLE {join_table_name} (\n'
                  f'  id INTEGER,\n'
                  f'  row_id INTEGER\n'
                  f')')
@@ -452,10 +459,22 @@ class SQLiteDatabase(Database):
             if isinstance(col, Tensor) and not isinstance(schema, TensorInfo):
                 self.schema[key] = schema = TensorInfo(
                     col.dtype,
+                    is_index=isinstance(col, Index),
                     is_edge_index=isinstance(col, EdgeIndex),
                 )
 
-            if isinstance(schema, TensorInfo) and schema.is_edge_index:
+            if isinstance(schema, TensorInfo) and schema.is_index:
+                assert isinstance(col, Index)
+
+                meta = torch.tensor([
+                    col.dim_size if col.dim_size is not None else -1,
+                    col.is_sorted,
+                ], dtype=torch.long)
+
+                out.append(meta.numpy().tobytes() +
+                           col.as_tensor().numpy().tobytes())
+
+            elif isinstance(schema, TensorInfo) and schema.is_edge_index:
                 assert isinstance(col, EdgeIndex)
 
                 num_rows, num_cols = col.sparse_size()
@@ -466,7 +485,8 @@ class SQLiteDatabase(Database):
                     col.is_undirected,
                 ], dtype=torch.long)
 
-                out.append(meta.numpy().tobytes() + col.numpy().tobytes())
+                out.append(meta.numpy().tobytes() +
+                           col.as_tensor().numpy().tobytes())
 
             elif isinstance(schema, TensorInfo):
                 assert isinstance(col, Tensor)
@@ -476,7 +496,9 @@ class SQLiteDatabase(Database):
                 out.append(col)
 
             else:
-                out.append(pickle.dumps(col))
+                buffer = io.BytesIO()
+                torch.save(col, buffer)
+                out.append(buffer.getvalue())
 
         return out
 
@@ -490,7 +512,23 @@ class SQLiteDatabase(Database):
         for i, (key, schema) in enumerate(self.schema.items()):
             value = row[i]
 
-            if isinstance(schema, TensorInfo) and schema.is_edge_index:
+            if isinstance(schema, TensorInfo) and schema.is_index:
+                meta = torch.frombuffer(value[:16], dtype=torch.long).tolist()
+                dim_size = meta[0] if meta[0] >= 0 else None
+                is_sorted = meta[1] > 0
+
+                if len(value) > 16:
+                    tensor = torch.frombuffer(value[16:], dtype=schema.dtype)
+                else:
+                    tensor = torch.empty(0, dtype=schema.dtype)
+
+                out_dict[key] = Index(
+                    tensor.view(*schema.size),
+                    dim_size=dim_size,
+                    is_sorted=is_sorted,
+                )
+
+            elif isinstance(schema, TensorInfo) and schema.is_edge_index:
                 meta = torch.frombuffer(value[:32], dtype=torch.long).tolist()
                 num_rows = meta[0] if meta[0] >= 0 else None
                 num_cols = meta[1] if meta[1] >= 0 else None
@@ -523,7 +561,10 @@ class SQLiteDatabase(Database):
                 out_dict[key] = value
 
             else:
-                out_dict[key] = pickle.loads(value)
+                out_dict[key] = torch.load(
+                    io.BytesIO(value),
+                    weights_only=False,
+                )
 
         # In case `0` exists as integer in the schema, this means that the
         # schema was passed as either a single entry or a tuple:
@@ -608,7 +649,12 @@ class RocksDatabase(Database):
         # Ensure that data is not a view of a larger tensor:
         if isinstance(row, Tensor):
             row = row.clone()
-        return pickle.dumps(row)
+        buffer = io.BytesIO()
+        torch.save(row, buffer)
+        return buffer.getvalue()
 
     def _deserialize(self, row: bytes) -> Any:
-        return pickle.loads(row)
+        return torch.load(
+            io.BytesIO(row),
+            weights_only=False,
+        )