Repository: pyg-team/pytorch_geometric
Pull request #8583 (2023-12-09 14:10:43)

`mypy` support in `torch_geometric.datasets` [2/n]

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8583?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`a8193ff`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/a8193ff2249c029ac4bc751132888bebb63084d4?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.76% compared to head [(`dbda9fb`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8583?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) 88.76%.

> :exclamation: Current head dbda9fb differs from pull request most recent head cc4f400. Consider uploading reports for the commit cc4f400 to get more accurate results


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #8583   +/-   ##
=======================================
  Coverage   88.76%   88.76%           
=======================================
  Files         479      479           
  Lines       29948    29948           
=======================================
  Hits        26583    26583           
  Misses       3365     3365           
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8583?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


Patch:
diff --git a/torch_geometric/data/dataset.py b/torch_geometric/data/dataset.py
index b862de8c1..86f05d2c8 100644
--- a/torch_geometric/data/dataset.py
+++ b/torch_geometric/data/dataset.py
@@ -16,6 +16,7 @@ from torch_geometric.data.data import BaseData
 from torch_geometric.io import fs
 
 IndexType = Union[slice, Tensor, np.ndarray, Sequence]
+MISSING = '???'
 
 
 class Dataset(torch.utils.data.Dataset, ABC):
@@ -49,24 +50,24 @@ class Dataset(torch.utils.data.Dataset, ABC):
             (default: :obj:`False`)
     """
     @property
-    def raw_file_names(self) -> Union[str, List[str], Tuple]:
+    def raw_file_names(self) -> Union[str, List[str], Tuple[str, ...]]:
         r"""The name of the files in the :obj:`self.raw_dir` folder that must
         be present in order to skip downloading.
         """
         raise NotImplementedError
 
     @property
-    def processed_file_names(self) -> Union[str, List[str], Tuple]:
+    def processed_file_names(self) -> Union[str, List[str], Tuple[str, ...]]:
         r"""The name of the files in the :obj:`self.processed_dir` folder that
         must be present in order to skip processing.
         """
         raise NotImplementedError
 
-    def download(self):
+    def download(self) -> None:
         r"""Downloads the dataset to the :obj:`self.raw_dir` folder."""
         raise NotImplementedError
 
-    def process(self):
+    def process(self) -> None:
         r"""Processes the dataset to the :obj:`self.processed_dir` folder."""
         raise NotImplementedError
 
@@ -88,13 +89,13 @@ class Dataset(torch.utils.data.Dataset, ABC):
         pre_filter: Optional[Callable] = None,
         log: bool = True,
         force_reload: bool = False,
-    ):
+    ) -> None:
         super().__init__()
 
         if isinstance(root, str):
             root = osp.expanduser(fs.normpath(root))
 
-        self.root = root or '???'
+        self.root = root or MISSING
         self.transform = transform
         self.pre_transform = pre_transform
         self.pre_filter = pre_filter
diff --git a/torch_geometric/data/in_memory_dataset.py b/torch_geometric/data/in_memory_dataset.py
index 9d7337152..781cbf556 100644
--- a/torch_geometric/data/in_memory_dataset.py
+++ b/torch_geometric/data/in_memory_dataset.py
@@ -2,14 +2,15 @@ import copy
 import os.path as osp
 import warnings
 from abc import ABC
-from collections.abc import Mapping, Sequence
 from typing import (
     Any,
     Callable,
     Dict,
     Iterable,
     List,
+    Mapping,
     Optional,
+    Sequence,
     Tuple,
     Type,
     Union,
@@ -61,11 +62,11 @@ class InMemoryDataset(Dataset, ABC):
             (default: :obj:`False`)
     """
     @property
-    def raw_file_names(self) -> Union[str, List[str], Tuple]:
+    def raw_file_names(self) -> Union[str, List[str], Tuple[str, ...]]:
         raise NotImplementedError
 
     @property
-    def processed_file_names(self) -> Union[str, List[str], Tuple]:
+    def processed_file_names(self) -> Union[str, List[str], Tuple[str, ...]]:
         raise NotImplementedError
 
     def __init__(
@@ -76,7 +77,7 @@ class InMemoryDataset(Dataset, ABC):
         pre_filter: Optional[Callable] = None,
         log: bool = True,
         force_reload: bool = False,
-    ):
+    ) -> None:
         super().__init__(root, transform, pre_transform, pre_filter, log,
                          force_reload)
 
@@ -120,12 +121,12 @@ class InMemoryDataset(Dataset, ABC):
         return data
 
     @classmethod
-    def save(cls, data_list: List[BaseData], path: str):
+    def save(cls, data_list: Sequence[BaseData], path: str) -> None:
         r"""Saves a list of data objects to the file path :obj:`path`."""
         data, slices = cls.collate(data_list)
         fs.torch_save((data.to_dict(), slices), path)
 
-    def load(self, path: str, data_cls: Type[BaseData] = Data):
+    def load(self, path: str, data_cls: Type[BaseData] = Data) -> None:
         r"""Loads the dataset from the file path :obj:`path`."""
         data, self.slices = fs.torch_load(path)
         if isinstance(data, dict):  # Backward compatibility.
@@ -134,7 +135,7 @@ class InMemoryDataset(Dataset, ABC):
 
     @staticmethod
     def collate(
-        data_list: List[BaseData],
+        data_list: Sequence[BaseData],
     ) -> Tuple[BaseData, Optional[Dict[str, Tensor]]]:
         r"""Collates a list of :class:`~torch_geometric.data.Data` or
         :class:`~torch_geometric.data.HeteroData` objects to the internal
diff --git a/torch_geometric/data/on_disk_dataset.py b/torch_geometric/data/on_disk_dataset.py
index 2976cc550..7abc1ea8d 100644
--- a/torch_geometric/data/on_disk_dataset.py
+++ b/torch_geometric/data/on_disk_dataset.py
@@ -1,5 +1,5 @@
 import os
-from typing import Any, Callable, Iterable, List, Optional, Union
+from typing import Any, Callable, Iterable, List, Optional, Sequence, Union
 
 from torch import Tensor
 
@@ -57,7 +57,7 @@ class OnDiskDataset(Dataset):
         backend: str = 'sqlite',
         schema: Schema = object,
         log: bool = True,
-    ):
+    ) -> None:
         if backend not in self.BACKENDS:
             raise ValueError(f"Database backend must be one of "
                              f"{set(self.BACKENDS.keys())} "
@@ -92,7 +92,7 @@ class OnDiskDataset(Dataset):
         self._numel = len(self._db)
         return self._db
 
-    def close(self):
+    def close(self) -> None:
         r"""Closes the connection to the underlying database."""
         if self._db is not None:
             self._db.close()
@@ -119,7 +119,7 @@ class OnDiskDataset(Dataset):
                                   f"needs to be overridden in case a "
                                   f"non-default schema was passed")
 
-    def append(self, data: BaseData):
+    def append(self, data: BaseData) -> None:
         r"""Appends the data object to the dataset."""
         index = len(self)
         self.db.insert(index, self.serialize(data))
@@ -127,9 +127,9 @@ class OnDiskDataset(Dataset):
 
     def extend(
         self,
-        data_list: List[BaseData],
+        data_list: Sequence[BaseData],
         batch_size: Optional[int] = None,
-    ):
+    ) -> None:
         r"""Extends the dataset by a list of data objects."""
         start = len(self)
         end = start + len(data_list)
diff --git a/torch_geometric/datasets/aqsol.py b/torch_geometric/datasets/aqsol.py
index 0cc2f2de4..b1c08d662 100644
--- a/torch_geometric/datasets/aqsol.py
+++ b/torch_geometric/datasets/aqsol.py
@@ -95,14 +95,14 @@ class AQSOL(InMemoryDataset):
     def processed_file_names(self) -> List[str]:
         return ['train.pt', 'val.pt', 'test.pt']
 
-    def download(self):
+    def download(self) -> None:
         fs.rm(self.raw_dir)
         path = download_url(self.url, self.root)
         extract_zip(path, self.root)
         os.rename(osp.join(self.root, 'asqol_graph_raw'), self.raw_dir)
         os.unlink(path)
 
-    def process(self):
+    def process(self) -> None:
         for raw_path, path in zip(self.raw_paths, self.processed_paths):
             with open(raw_path, 'rb') as f:
                 graphs = pickle.load(f)
diff --git a/torch_geometric/datasets/ba2motif_dataset.py b/torch_geometric/datasets/ba2motif_dataset.py
index e665c8502..5a66dce94 100644
--- a/torch_geometric/datasets/ba2motif_dataset.py
+++ b/torch_geometric/datasets/ba2motif_dataset.py
@@ -91,6 +91,7 @@ class BA2MotifDataset(InMemoryDataset):
                          force_reload=force_reload)
         self.load(self.processed_paths[0])
 
+    @property
     def raw_file_names(self) -> str:
         return self.filename
 
diff --git a/torch_geometric/datasets/bitcoin_otc.py b/torch_geometric/datasets/bitcoin_otc.py
index 49c9417c0..4b2b3ba3e 100644
--- a/torch_geometric/datasets/bitcoin_otc.py
+++ b/torch_geometric/datasets/bitcoin_otc.py
@@ -77,7 +77,9 @@ class BitcoinOTC(InMemoryDataset):
 
     @property
     def num_nodes(self) -> int:
-        return self._data.edge_index.max().item() + 1
+        assert isinstance(self._data, Data)
+        assert self._data.edge_index is not None
+        return int(self._data.edge_index.max()) + 1
 
     def download(self):
         path = download_url(self.url, self.raw_dir)
diff --git a/torch_geometric/datasets/dbp15k.py b/torch_geometric/datasets/dbp15k.py
index 16388644c..daed8b2d9 100644
--- a/torch_geometric/datasets/dbp15k.py
+++ b/torch_geometric/datasets/dbp15k.py
@@ -122,14 +122,14 @@ class DBP15K(InMemoryDataset):
                 x_dict[int(info[0])] = torch.stack(hs, dim=0)
 
         idx = torch.tensor(list(x_dict.keys()))
-        assoc = torch.full((idx.max().item() + 1, ), -1, dtype=torch.long)
+        assoc = torch.full((int(idx.max()) + 1, ), -1, dtype=torch.long)
         assoc[idx] = torch.arange(idx.size(0))
 
         subj, obj = assoc[subj], assoc[obj]
         edge_index = torch.stack([subj, obj], dim=0)
         edge_index, rel = sort_edge_index(edge_index, rel)
 
-        xs = [None for _ in range(idx.size(0))]
+        xs = list(x_dict.values())
         for i in x_dict.keys():
             xs[assoc[i]] = x_dict[i]
         x = torch.nn.utils.rnn.pad_sequence(xs, batch_first=True)
diff --git a/torch_geometric/datasets/entities.py b/torch_geometric/datasets/entities.py
index c6edf5c01..5dc09cb42 100644
--- a/torch_geometric/datasets/entities.py
+++ b/torch_geometric/datasets/entities.py
@@ -104,11 +104,11 @@ class Entities(InMemoryDataset):
 
     @property
     def num_relations(self) -> int:
-        return self._data.edge_type.max().item() + 1
+        return int(self._data.edge_type.max()) + 1  # type: ignore
 
     @property
     def num_classes(self) -> int:
-        return self._data.train_y.max().item() + 1
+        return int(self._data.train_y.max()) + 1  # type: ignore
 
     @property
     def raw_file_names(self) -> List[str]:
diff --git a/torch_geometric/datasets/gdelt.py b/torch_geometric/datasets/gdelt.py
index 45e3ac2f7..9b1ac7cfb 100644
--- a/torch_geometric/datasets/gdelt.py
+++ b/torch_geometric/datasets/gdelt.py
@@ -1,6 +1,7 @@
 from typing import Callable, List, Optional
 
 import torch
+from torch import Tensor
 
 from torch_geometric.data import download_url
 from torch_geometric.datasets.icews import EventDataset
@@ -72,7 +73,7 @@ class GDELT(EventDataset):
         for filename in self.raw_file_names:
             download_url(f'{self.url}/{filename}', self.raw_dir)
 
-    def process_events(self) -> torch.Tensor:
+    def process_events(self) -> Tensor:
         events = []
         for path in self.raw_paths:
             data = read_txt_array(path, sep='\t', end=4, dtype=torch.long)
@@ -82,7 +83,7 @@ class GDELT(EventDataset):
 
     def process(self):
         s = self.splits
-        data_list = super().process()
+        data_list = self._process_data_list()
         self.save(data_list[s[0]:s[1]], self.processed_paths[0])
         self.save(data_list[s[1]:s[2]], self.processed_paths[1])
         self.save(data_list[s[2]:s[3]], self.processed_paths[2])
diff --git a/torch_geometric/datasets/graph_generator/base.py b/torch_geometric/datasets/graph_generator/base.py
index ad54dddca..7bf8a41dc 100644
--- a/torch_geometric/datasets/graph_generator/base.py
+++ b/torch_geometric/datasets/graph_generator/base.py
@@ -14,13 +14,13 @@ class GraphGenerator(ABC):
 
     @staticmethod
     def resolve(query: Union[Any, str], *args, **kwargs):
-        import torch_geometric.datasets.graph_generator as graph_generators
+        import torch_geometric.datasets.graph_generator as _graph_generators
         graph_generators = [
-            gen for gen in vars(graph_generators).values()
+            gen for gen in vars(_graph_generators).values()
             if isinstance(gen, type) and issubclass(gen, GraphGenerator)
         ]
-        return resolver(graph_generators, {}, query, base_cls=GraphGenerator,
-                        base_cls_repr='Graph', *args, **kwargs)
+        return resolver(graph_generators, {}, query, GraphGenerator, 'Graph',
+                        *args, **kwargs)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}()'
diff --git a/torch_geometric/datasets/icews.py b/torch_geometric/datasets/icews.py
index 38ba24698..e5d128983 100644
--- a/torch_geometric/datasets/icews.py
+++ b/torch_geometric/datasets/icews.py
@@ -1,6 +1,7 @@
 from typing import Callable, List, Optional
 
 import torch
+from torch import Tensor
 
 from torch_geometric.data import Data, InMemoryDataset, download_url
 from torch_geometric.io import read_txt_array
@@ -26,10 +27,10 @@ class EventDataset(InMemoryDataset):
     def num_rels(self) -> int:
         raise NotImplementedError
 
-    def process_events(self) -> int:
+    def process_events(self) -> Tensor:
         raise NotImplementedError
 
-    def process(self) -> List[Data]:
+    def _process_data_list(self) -> List[Data]:
         events = self.process_events()
         events = events - events.min(dim=0, keepdim=True)[0]
 
@@ -110,7 +111,7 @@ class ICEWS18(EventDataset):
         for filename in self.raw_file_names:
             download_url(f'{self.url}/{filename}', self.raw_dir)
 
-    def process_events(self) -> torch.Tensor:
+    def process_events(self) -> Tensor:
         events = []
         for path in self.raw_paths:
             data = read_txt_array(path, sep='\t', end=4, dtype=torch.long)
@@ -120,7 +121,7 @@ class ICEWS18(EventDataset):
 
     def process(self):
         s = self.splits
-        data_list = super().process()
+        data_list = self._process_data_list()
         self.save(data_list[s[0]:s[1]], self.processed_paths[0])
         self.save(data_list[s[1]:s[2]], self.processed_paths[1])
         self.save(data_list[s[2]:s[3]], self.processed_paths[2])
diff --git a/torch_geometric/datasets/infection_dataset.py b/torch_geometric/datasets/infection_dataset.py
index 83eea101c..c05456dd4 100644
--- a/torch_geometric/datasets/infection_dataset.py
+++ b/torch_geometric/datasets/infection_dataset.py
@@ -86,6 +86,8 @@ class InfectionDataset(InMemoryDataset):
         if num_graphs is None and isinstance(max_path_length, list):
             num_graphs = len(max_path_length)
 
+        assert num_graphs is not None
+
         self.graph_generator = GraphGenerator.resolve(
             graph_generator,
             **(graph_generator_kwargs or {}),
diff --git a/torch_geometric/datasets/md17.py b/torch_geometric/datasets/md17.py
index 65e67a1ab..9ada76970 100644
--- a/torch_geometric/datasets/md17.py
+++ b/torch_geometric/datasets/md17.py
@@ -1,6 +1,6 @@
 import os
 import os.path as osp
-from typing import Callable, List, Optional
+from typing import Callable, List, Optional, Union
 
 import numpy as np
 import torch
@@ -368,6 +368,7 @@ class MD17(InMemoryDataset):
         self.load(self.processed_paths[idx])
 
     def mean(self) -> float:
+        assert isinstance(self._data, Data)
         return float(self._data.energy.mean())
 
     @property
@@ -381,12 +382,12 @@ class MD17(InMemoryDataset):
         return osp.join(self.root, self.name, 'processed')
 
     @property
-    def raw_file_names(self) -> str:
+    def raw_file_names(self) -> Union[str, List[str]]:
         name = self.file_names[self.name]
         if self.revised:
             return osp.join('rmd17', 'npz_data', name)
         elif self.ccsd:
-            return name[:-4] + '-train.npz', name[:-4] + '-test.npz'
+            return [name[:-4] + '-train.npz', name[:-4] + '-test.npz']
         return name
 
     @property
diff --git a/torch_geometric/datasets/modelnet.py b/torch_geometric/datasets/modelnet.py
index 3c5a13fa7..19d4ba8aa 100644
--- a/torch_geometric/datasets/modelnet.py
+++ b/torch_geometric/datasets/modelnet.py
@@ -1,10 +1,9 @@
 import glob
 import os
 import os.path as osp
-from typing import Callable, Dict, List, Optional, Tuple
+from typing import Callable, List, Optional
 
 import torch
-from torch import Tensor
 
 from torch_geometric.data import (
     Data,
@@ -129,7 +128,7 @@ class ModelNet(InMemoryDataset):
         self.save(self.process_set('train'), self.processed_paths[0])
         self.save(self.process_set('test'), self.processed_paths[1])
 
-    def process_set(self, dataset: str) -> Tuple[Data, Dict[str, Tensor]]:
+    def process_set(self, dataset: str) -> List[Data]:
         categories = glob.glob(osp.join(self.raw_dir, '*', ''))
         categories = sorted([x.split(os.sep)[-2] for x in categories])
 
diff --git a/torch_geometric/datasets/motif_generator/base.py b/torch_geometric/datasets/motif_generator/base.py
index a6057a85a..bf5b9b429 100644
--- a/torch_geometric/datasets/motif_generator/base.py
+++ b/torch_geometric/datasets/motif_generator/base.py
@@ -14,13 +14,13 @@ class MotifGenerator(ABC):
 
     @staticmethod
     def resolve(query: Union[Any, str], *args, **kwargs):
-        import torch_geometric.datasets.motif_generator as motif_generators
+        import torch_geometric.datasets.motif_generator as _motif_generators
         motif_generators = [
-            gen for gen in vars(motif_generators).values()
+            gen for gen in vars(_motif_generators).values()
             if isinstance(gen, type) and issubclass(gen, MotifGenerator)
         ]
-        return resolver(motif_generators, {}, query, base_cls=MotifGenerator,
-                        base_cls_repr='Motif', *args, **kwargs)
+        return resolver(motif_generators, {}, query, MotifGenerator, 'Motif',
+                        *args, **kwargs)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}()'
diff --git a/torch_geometric/datasets/motif_generator/custom.py b/torch_geometric/datasets/motif_generator/custom.py
index 015561e3b..ecd335468 100644
--- a/torch_geometric/datasets/motif_generator/custom.py
+++ b/torch_geometric/datasets/motif_generator/custom.py
@@ -1,4 +1,4 @@
-from typing import Any
+from typing import Any, Optional
 
 from torch_geometric.data import Data
 from torch_geometric.datasets.motif_generator import MotifGenerator
@@ -16,7 +16,8 @@ class CustomMotif(MotifGenerator):
     def __init__(self, structure: Any):
         super().__init__()
 
-        self.structure = None
+        self.structure: Optional[Data] = None
+
         if isinstance(structure, Data):
             self.structure = structure
         else:
@@ -33,4 +34,5 @@ class CustomMotif(MotifGenerator):
                              f"(got {type(structure)})")
 
     def __call__(self) -> Data:
+        assert isinstance(self.structure, Data)
         return self.structure
diff --git a/torch_geometric/datasets/ogb_mag.py b/torch_geometric/datasets/ogb_mag.py
index 5db4d7b6b..63843bf67 100644
--- a/torch_geometric/datasets/ogb_mag.py
+++ b/torch_geometric/datasets/ogb_mag.py
@@ -72,6 +72,7 @@ class OGB_MAG(InMemoryDataset):
 
     @property
     def num_classes(self) -> int:
+        assert isinstance(self._data, HeteroData)
         return int(self._data['paper'].y.max()) + 1
 
     @property
diff --git a/torch_geometric/datasets/pcqm4m.py b/torch_geometric/datasets/pcqm4m.py
index a9efa3b3f..8543f4f88 100644
--- a/torch_geometric/datasets/pcqm4m.py
+++ b/torch_geometric/datasets/pcqm4m.py
@@ -6,6 +6,7 @@ import torch
 from tqdm import tqdm
 
 from torch_geometric.data import Data, OnDiskDataset, download_url, extract_zip
+from torch_geometric.data.data import BaseData
 from torch_geometric.utils import from_smiles
 
 
@@ -80,7 +81,7 @@ class PCQM4Mv2(OnDiskDataset):
         extract_zip(path, self.raw_dir)
         os.unlink(path)
 
-    def process(self):
+    def process(self) -> None:
         import pandas as pd
 
         df = pd.read_csv(self.raw_paths[0])
@@ -96,7 +97,8 @@ class PCQM4Mv2(OnDiskDataset):
                 self.extend(data_list)
                 data_list = []
 
-    def serialize(self, data: Data) -> Dict[str, Any]:
+    def serialize(self, data: BaseData) -> Dict[str, Any]:
+        assert isinstance(data, Data)
         return dict(
             x=data.x,
             edge_index=data.edge_index,
diff --git a/torch_geometric/datasets/ppi.py b/torch_geometric/datasets/ppi.py
index 31308a522..a3eb36bfc 100644
--- a/torch_geometric/datasets/ppi.py
+++ b/torch_geometric/datasets/ppi.py
@@ -92,7 +92,7 @@ class PPI(InMemoryDataset):
         return [f'{split}_{name}' for split, name in product(splits, files)]
 
     @property
-    def processed_file_names(self) -> str:
+    def processed_file_names(self) -> List[str]:
         return ['train.pt', 'val.pt', 'test.pt']
 
     def download(self):
diff --git a/torch_geometric/datasets/shapenet.py b/torch_geometric/datasets/shapenet.py
index dd8b799b7..3c9425e92 100644
--- a/torch_geometric/datasets/shapenet.py
+++ b/torch_geometric/datasets/shapenet.py
@@ -152,6 +152,8 @@ class ShapeNet(InMemoryDataset):
                               'train, val, trainval or test'))
 
         self.load(path)
+
+        assert isinstance(self._data, Data)
         self._data.x = self._data.x if include_normals else None
 
         self.y_mask = torch.zeros((len(self.seg_classes.keys()), 50),
@@ -168,7 +170,7 @@ class ShapeNet(InMemoryDataset):
         return list(self.category_ids.values()) + ['train_test_split']
 
     @property
-    def processed_file_names(self) -> str:
+    def processed_file_names(self) -> List[str]:
         cats = '_'.join([cat[:3].lower() for cat in self.categories])
         return [
             osp.join(f'{cats}_{split}.pt')


Test patch:
diff --git a/test/nn/models/test_re_net.py b/test/nn/models/test_re_net.py
index 0a5641e19..e0761e81a 100644
--- a/test/nn/models/test_re_net.py
+++ b/test/nn/models/test_re_net.py
@@ -34,7 +34,7 @@ class MyTestEventDataset(EventDataset):
         return torch.stack([sub, rel, obj, t], dim=1)
 
     def process(self):
-        data_list = super().process()
+        data_list = self._process_data_list()
         self.save(data_list, self.processed_paths[0])