Repository: pyg-team/pytorch_geometric
Pull request #9234 (2024-04-24 12:14:39)

Fixed node mapping in `RCDD` dataset
Fix https://github.com/pyg-team/pytorch_geometric/issues/9212

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9234?dropdown=coverage&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:
> Project coverage is 82.52%. Comparing base [(`ed17034`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/ed170342eb2b174fd16b910c735758edbd4e78fd?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`8fc42b4`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9234?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).


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


```diff
@@            Coverage Diff             @@
##           master    #9234      +/-   ##
==========================================
- Coverage   88.48%   82.52%   -5.97%     
==========================================
  Files         470      470              
  Lines       30189    30189              
==========================================
- Hits        26714    24913    -1801     
- Misses       3475     5276    +1801     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9234?dropdown=coverage&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/.github/actions/setup/action.yml b/.github/actions/setup/action.yml
index 468b6922e..2c5a5e916 100644
--- a/.github/actions/setup/action.yml
+++ b/.github/actions/setup/action.yml
@@ -6,13 +6,13 @@ inputs:
     default: '3.8'
   torch-version:
     required: false
-    default: '2.2.0'
+    default: '2.3.0'
   cuda-version:
     required: false
     default: cpu
   torchvision-version:
     required: false
-    default: '0.17.0'
+    default: '0.18.0'
   full_install:
     required: false
     default: true
@@ -49,7 +49,7 @@ runs:
       shell: bash
 
     - name: Install pyg-lib  # pyg-lib is currently not available on Windows.
-      if: ${{ inputs.torch-version != 'nightly' && runner.os != 'Windows' }}
+      if: ${{ inputs.torch-version != 'nightly' && inputs.torch-version != '2.3.0' && runner.os != 'Windows' }}
       run: |
         pip uninstall -y pyg-lib
         pip install --no-index pyg-lib -f https://data.pyg.org/whl/nightly/torch-${{ inputs.torch-version }}+${{ inputs.cuda-version }}.html
@@ -74,7 +74,7 @@ runs:
       shell: bash
 
     - name: Install extension packages
-      if: ${{ inputs.full_install == 'true' && inputs.torch-version != 'nightly' && runner.os != 'macOS' }}
+      if: ${{ inputs.full_install == 'true' && inputs.torch-version != 'nightly' && inputs.torch-version != '2.3.0' && runner.os != 'macOS' }}
       run: |
         pip install scipy
         pip install --no-index --upgrade torch-scatter torch-sparse torch-cluster torch-spline-conv -f https://data.pyg.org/whl/torch-${{ inputs.torch-version }}+${{ inputs.cuda-version }}.html
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc4043841..4632892d5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for PyTorch 2.3 ([#9240](https://github.com/pyg-team/pytorch_geometric/pull/9240))
 - Added support for `EdgeIndex` in `message_and_aggregate` ([#9131](https://github.com/pyg-team/pytorch_geometric/pull/9131))
 - Added `CornellTemporalHyperGraphDataset` ([#9090](https://github.com/pyg-team/pytorch_geometric/pull/9090))
 - Added support for cuGraph data loading and `GAT` in single node Papers100m examples ([#8173](https://github.com/pyg-team/pytorch_geometric/pull/8173))
@@ -38,6 +39,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fix node mapping bug in `RCDD` dataset ([#9234](https://github.com/pyg-team/pytorch_geometric/pull/9234))
 - Fixed incorrect treatment of `edge_label` and `edge_label_index` in `ToSparseTensor` transform ([#9199](https://github.com/pyg-team/pytorch_geometric/pull/9199))
 - Fixed `EgoData` processing in `SnapDataset` in case filenames are unsorted ([#9195](https://github.com/pyg-team/pytorch_geometric/pull/9195))
 - Fixed empty graph and isolated node handling in `to_dgl` ([#9188](https://github.com/pyg-team/pytorch_geometric/pull/9188))
diff --git a/README.md b/README.md
index 7dcb95b85..6481c7e16 100644
--- a/README.md
+++ b/README.md
@@ -16,7 +16,7 @@ ______________________________________________________________________
 **PyG** *(PyTorch Geometric)* is a library built upon [PyTorch](https://pytorch.org/) to easily write and train Graph Neural Networks (GNNs) for a wide range of applications related to structured data.
 
 It consists of various methods for deep learning on graphs and other irregular structures, also known as *[geometric deep learning](http://geometricdeeplearning.com/)*, from a variety of published papers.
-In addition, it consists of easy-to-use mini-batch loaders for operating on many small and single giant graphs, [multi GPU-support](https://github.com/pyg-team/pytorch_geometric/tree/master/examples/multi_gpu), [`torch.compile`](https://pytorch-geometric.readthedocs.io/en/latest/advanced/compile.html) support, [`DataPipe`](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/datapipe.py) support, a large number of common benchmark datasets (based on simple interfaces to create your own), the [GraphGym](https://pytorch-geometric.readthedocs.io/en/latest/advanced/graphgym.html) experiment manager, and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.
+In addition, it consists of easy-to-use mini-batch loaders for operating on many small and single giant graphs, [multi GPU-support](https://github.com/pyg-team/pytorch_geometric/tree/master/examples/multi_gpu), [`torch.compile`](https://pytorch-geometric.readthedocs.io/en/latest/advanced/compile.html) support, [`DataPipe`](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/datapipe.py) support, a large number of common benchmark datasets (based on simple interfaces to create your own), and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.
 
 **[Click here to join our Slack community!][slack-url]**
 
@@ -47,7 +47,6 @@ Whether you are a machine learning researcher or first-time user of machine lear
   Making modifications to existing models or creating new architectures is simple, thanks to its easy-to-use message passing API, and a variety of operators and utility functions.
 - **Large-scale real-world GNN models**:
   We focus on the need of GNN applications in challenging real-world scenarios, and support learning on diverse types of graphs, including but not limited to: scalable GNNs for graphs with millions of nodes; dynamic GNNs for node predictions over time; heterogeneous GNNs with multiple node types and edge types.
-- **GraphGym integration**: GraphGym lets users easily reproduce GNN experiments, is able to launch and analyze thousands of different GNN configurations, and is customizable by registering new modules to a GNN learning pipeline.
 
 ## Quick Tour for New Users
 
@@ -139,20 +138,6 @@ class EdgeConv(MessagePassing):
         return self.mlp(edge_features)  # shape [num_edges, out_channels]
 ```
 
-### Manage experiments with GraphGym
-
-GraphGym allows you to manage and launch GNN experiments, using a highly modularized pipeline (see [here](https://pytorch-geometric.readthedocs.io/en/latest/advanced/graphgym.html) for the accompanying tutorial).
-
-```
-git clone https://github.com/pyg-team/pytorch_geometric.git
-cd pytorch_geometric/graphgym
-bash run_single.sh  # run a single GNN experiment (node/edge/graph-level)
-bash run_batch.sh   # run a batch of GNN experiments, using differnt GNN designs/datasets/tasks
-```
-
-Users are highly encouraged to check out the [documentation](https://pytorch-geometric.readthedocs.io/en/latest), which contains additional tutorials on the essential functionalities of PyG, including data handling, creation of datasets and a full list of implemented methods, transforms, and datasets.
-For a quick start, check out our [examples](https://github.com/pyg-team/pytorch_geometric/tree/master/examples) in `examples/`.
-
 ## Architecture Overview
 
 PyG provides a multi-layer framework that enables users to build Graph Neural Network solutions on both low and high levels.
diff --git a/docs/source/index.rst b/docs/source/index.rst
index dcef5e18e..1c67eedde 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -6,7 +6,7 @@ PyG Documentation
 :pyg:`null` **PyG** *(PyTorch Geometric)* is a library built upon :pytorch:`null` `PyTorch <https://pytorch.org>`_ to easily write and train Graph Neural Networks (GNNs) for a wide range of applications related to structured data.
 
 It consists of various methods for deep learning on graphs and other irregular structures, also known as `geometric deep learning <http://geometricdeeplearning.com/>`_, from a variety of published papers.
-In addition, it consists of easy-to-use mini-batch loaders for operating on many small and single giant graphs, `multi GPU-support <https://github.com/pyg-team/pytorch_geometric/tree/master/examples/multi_gpu>`_, `torch.compile <https://pytorch-geometric.readthedocs.io/en/latest/advanced/compile.html>`_ support, `DataPipe <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/datapipe.py>`_ support, a large number of common benchmark datasets (based on simple interfaces to create your own), the `GraphGym <advanced/graphgym.html>`__ experiment manager, and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.
+In addition, it consists of easy-to-use mini-batch loaders for operating on many small and single giant graphs, `multi GPU-support <https://github.com/pyg-team/pytorch_geometric/tree/master/examples/multi_gpu>`_, `torch.compile <https://pytorch-geometric.readthedocs.io/en/latest/advanced/compile.html>`_ support, `DataPipe <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/datapipe.py>`_ support, a large number of common benchmark datasets (based on simple interfaces to create your own), and helpful transforms, both for learning on arbitrary graphs as well as on 3D meshes or point clouds.
 
 .. slack_button::
 
diff --git a/examples/multi_gpu/model_parallel.py b/examples/multi_gpu/model_parallel.py
index f939fdc03..b45961955 100644
--- a/examples/multi_gpu/model_parallel.py
+++ b/examples/multi_gpu/model_parallel.py
@@ -19,6 +19,9 @@ data = dataset[0].to('cuda:0')
 class GCN(torch.nn.Module):
     def __init__(self, in_channels, out_channels, device1, device2):
         super().__init__()
+        self.device1 = device1
+        self.device2 = device2
+
         self.conv1 = GCNConv(in_channels, 16).to(device1)
         self.conv2 = GCNConv(16, out_channels).to(device2)
 
@@ -54,7 +57,7 @@ def train():
 @torch.no_grad()
 def test():
     model.eval()
-    out = model(data.x, data.edge_index, data.edge_attr)
+    out = model(data.x, data.edge_index)
     pred = out.argmax(dim=-1).to('cuda:0')
 
     accs = []
diff --git a/torch_geometric/datasets/karate.py b/torch_geometric/datasets/karate.py
index 7da24dd3e..77473a69c 100644
--- a/torch_geometric/datasets/karate.py
+++ b/torch_geometric/datasets/karate.py
@@ -8,8 +8,9 @@ from torch_geometric.data import Data, InMemoryDataset
 class KarateClub(InMemoryDataset):
     r"""Zachary's karate club network from the `"An Information Flow Model for
     Conflict and Fission in Small Groups"
-    <http://www1.ind.ku.dk/complexLearning/zachary1977.pdf>`_ paper, containing
-    34 nodes, connected by 156 (undirected and unweighted) edges.
+    <https://www.journals.uchicago.edu/doi/abs/10.1086/jar.33.4.3629752>`_
+    paper, containing 34 nodes,
+    connected by 156 (undirected and unweighted) edges.
     Every node is labeled by one of four classes obtained via modularity-based
     clustering, following the `"Semi-supervised Classification with Graph
     Convolutional Networks" <https://arxiv.org/abs/1609.02907>`_ paper.
diff --git a/torch_geometric/datasets/rcdd.py b/torch_geometric/datasets/rcdd.py
index c83713229..33a5197ad 100644
--- a/torch_geometric/datasets/rcdd.py
+++ b/torch_geometric/datasets/rcdd.py
@@ -85,13 +85,13 @@ class RCDD(InMemoryDataset):
         mapping = torch.empty(len(node_df), dtype=torch.long)
         for node_type in node_df['node_type'].unique():
             mask = node_df['node_type'] == node_type
-            mask = torch.from_numpy(mask.values)
-            num_nodes = int(mask.sum())
-            mapping[mask] = torch.arange(num_nodes)
+            node_id = torch.from_numpy(node_df['node_id'][mask].values)
+            num_nodes = mask.sum()
+            mapping[node_id] = torch.arange(num_nodes)
             data[node_type].num_nodes = num_nodes
             x = np.vstack([
                 np.asarray(f.split(':'), dtype=np.float32)
-                for f in node_df['node_feat'][mask.numpy()]
+                for f in node_df['node_feat'][mask]
             ])
             data[node_type].x = torch.from_numpy(x)
 
diff --git a/torch_geometric/datasets/s3dis.py b/torch_geometric/datasets/s3dis.py
index faef077b9..6f574f66a 100644
--- a/torch_geometric/datasets/s3dis.py
+++ b/torch_geometric/datasets/s3dis.py
@@ -3,6 +3,7 @@ import os.path as osp
 from typing import Callable, List, Optional
 
 import torch
+from torch import Tensor
 
 from torch_geometric.data import (
     Data,
@@ -91,7 +92,8 @@ class S3DIS(InMemoryDataset):
         with open(self.raw_paths[1], 'r') as f:
             rooms = f.read().split('\n')[:-1]
 
-        xs, ys = [], []
+        xs: List[Tensor] = []
+        ys: List[Tensor] = []
         for filename in filenames:
             h5 = h5py.File(osp.join(self.raw_dir, filename))
             xs += torch.from_numpy(h5['data'][:]).unbind(0)
diff --git a/torch_geometric/edge_index.py b/torch_geometric/edge_index.py
index b5d3721dc..b2478d75f 100644
--- a/torch_geometric/edge_index.py
+++ b/torch_geometric/edge_index.py
@@ -722,6 +722,25 @@ class EdgeIndex(Tensor):
 
     # Methods #################################################################
 
+    def share_memory_(self) -> 'EdgeIndex':
+        self._data.share_memory_()
+        if self._indptr is not None:
+            self._indptr.share_memory_()
+        if self._T_perm is not None:
+            self._T_perm.share_memory_()
+        if self._T_index[0] is not None:
+            self._T_index[0].share_memory_()
+        if self._T_index[1] is not None:
+            self._T_index[1].share_memory_()
+        if self._T_indptr is not None:
+            self._T_indptr.share_memory_()
+        if self._value is not None:
+            self._value.share_memory_()
+        return self
+
+    def is_shared(self) -> bool:
+        return self._data.is_shared()
+
     def as_tensor(self) -> Tensor:
         r"""Zero-copies the :class:`EdgeIndex` representation back to a
         :class:`torch.Tensor` representation.
@@ -1111,6 +1130,8 @@ class EdgeIndex(Tensor):
             edge_index._indptr = colptr
             return edge_index
 
+    # PyTorch/Python builtins #################################################
+
     def __tensor_flatten__(self) -> Tuple[List[str], Tuple[Any, ...]]:
         attrs = ['_data']
         if self._indptr is not None:
@@ -1208,7 +1229,9 @@ class EdgeIndex(Tensor):
         return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
                                                indent, force_newline=False)
 
-    def shallow_copy(self) -> 'EdgeIndex':
+    # Helpers #################################################################
+
+    def _shallow_copy(self) -> 'EdgeIndex':
         out = EdgeIndex(self._data)
         out._sparse_size = self._sparse_size
         out._sort_order = self._sort_order
@@ -1221,25 +1244,18 @@ class EdgeIndex(Tensor):
         out._cat_metadata = self._cat_metadata
         return out
 
-    def share_memory_(self) -> 'EdgeIndex':
-        self._data.share_memory_()
-        if self._indptr is not None:
-            self._indptr.share_memory_()
-        if self._T_perm is not None:
-            self._T_perm.share_memory_()
-        if self._T_index[0] is not None:
-            self._T_index[0].share_memory_()
-        if self._T_index[1] is not None:
-            self._T_index[1].share_memory_()
-        if self._T_indptr is not None:
-            self._T_indptr.share_memory_()
-        if self._value is not None:
-            self._value.share_memory_()
+    def _clear_metadata(self) -> 'EdgeIndex':
+        self._sparse_size = (None, None)
+        self._sort_order = None
+        self._is_undirected = False
+        self._indptr = None
+        self._T_perm = None
+        self._T_index = (None, None)
+        self._T_indptr = None
+        self._value = None
+        self._cat_metadata = None
         return self
 
-    def is_shared(self) -> bool:
-        return self._data.is_shared()
-
 
 class SortReturnType(NamedTuple):
     values: EdgeIndex
@@ -1435,7 +1451,7 @@ def _slice(
 
     if ((start is None or start <= 0)
             and (end is None or end > input.size(dim)) and step == 1):
-        return input.shallow_copy()  # No-op.
+        return input._shallow_copy()  # No-op.
 
     out = aten.slice.Tensor(input._data, dim, start, end, step)
 
@@ -1539,19 +1555,52 @@ def _add(
     return out
 
 
-# @implements(aten.add_.Tensor)
-# def add_(
-#     input: EdgeIndex,
-#     other: Union[int, Tensor, EdgeIndex],
-#     *,
-#     alpha: int = 1,
-# ) -> Tensor:
-#     aten.add_.Tensor(
-#         input._data,
-#         other._data if isinstance(other, EdgeIndex) else other,
-#         alpha=alpha,
-#     )
-#     return input
+@implements(aten.add_.Tensor)
+def add_(
+    input: EdgeIndex,
+    other: Union[int, Tensor, EdgeIndex],
+    *,
+    alpha: int = 1,
+) -> EdgeIndex:
+
+    sparse_size = input._sparse_size
+    sort_order = input._sort_order
+    is_undirected = input._is_undirected
+    T_perm = input._T_perm
+    input._clear_metadata()
+
+    aten.add_.Tensor(
+        input._data,
+        other._data if isinstance(other, EdgeIndex) else other,
+        alpha=alpha,
+    )
+
+    if isinstance(other, Tensor) and other.numel() <= 1:
+        other = int(other)
+
+    if isinstance(other, int):
+        size = maybe_add(sparse_size, other, alpha)
+        assert len(size) == 2
+        input._sparse_size = size
+        input._sort_order = sort_order
+        input._is_undirected = is_undirected
+        input._T_perm = T_perm
+
+    elif isinstance(other, Tensor) and other.size() == (2, 1):
+        size = maybe_add(sparse_size, other.view(-1).tolist(), alpha)
+        assert len(size) == 2
+        input._sparse_size = size
+        input._sort_order = sort_order
+        if torch.equal(other[0], other[1]):
+            input._is_undirected = is_undirected
+        input._T_perm = T_perm
+
+    elif isinstance(other, EdgeIndex):
+        size = maybe_add(sparse_size, other._sparse_size, alpha)
+        assert len(size) == 2
+        input._sparse_size = size
+
+    return input
 
 
 @implements(aten.sub.Tensor)
@@ -1598,19 +1647,48 @@ def _sub(
     return out
 
 
-# @implements(aten.sub_.Tensor)
-# def sub_(
-#     input: EdgeIndex,
-#     other: Union[int, Tensor, EdgeIndex],
-#     *,
-#     alpha: int = 1,
-# ) -> EdgeIndex:
-#     aten.sub_.Tensor(
-#         input._data,
-#         other._data if isinstance(other, EdgeIndex) else other,
-#         alpha=alpha,
-#     )
-#     return input
+@implements(aten.sub_.Tensor)
+def sub_(
+    input: EdgeIndex,
+    other: Union[int, Tensor, EdgeIndex],
+    *,
+    alpha: int = 1,
+) -> EdgeIndex:
+
+    sparse_size = input._sparse_size
+    sort_order = input._sort_order
+    is_undirected = input._is_undirected
+    T_perm = input._T_perm
+    input._clear_metadata()
+
+    aten.sub_.Tensor(
+        input._data,
+        other._data if isinstance(other, EdgeIndex) else other,
+        alpha=alpha,
+    )
+
+    if isinstance(other, Tensor) and other.numel() <= 1:
+        other = int(other)
+
+    if isinstance(other, int):
+        size = maybe_sub(sparse_size, other, alpha)
+        assert len(size) == 2
+        input._sparse_size = size
+        input._sort_order = sort_order
+        input._is_undirected = is_undirected
+        input._T_perm = T_perm
+
+    elif isinstance(other, Tensor) and other.size() == (2, 1):
+        size = maybe_sub(sparse_size, other.view(-1).tolist(), alpha)
+        assert len(size) == 2
+        input._sparse_size = size
+        input._sort_order = sort_order
+        if torch.equal(other[0], other[1]):
+            input._is_undirected = is_undirected
+        input._T_perm = T_perm
+
+    return input
+
 
 # Sparse-Dense Matrix Multiplication ##########################################
 
diff --git a/torch_geometric/io/tu.py b/torch_geometric/io/tu.py
index dc712e2f7..adb5a0d77 100644
--- a/torch_geometric/io/tu.py
+++ b/torch_geometric/io/tu.py
@@ -37,7 +37,7 @@ def read_tu_data(
         if node_label.dim() == 1:
             node_label = node_label.unsqueeze(-1)
         node_label = node_label - node_label.min(dim=0)[0]
-        node_labels = node_label.unbind(dim=-1)
+        node_labels = list(node_label.unbind(dim=-1))
         node_labels = [one_hot(x) for x in node_labels]
         if len(node_labels) == 1:
             node_label = node_labels[0]
@@ -56,7 +56,7 @@ def read_tu_data(
         if edge_label.dim() == 1:
             edge_label = edge_label.unsqueeze(-1)
         edge_label = edge_label - edge_label.min(dim=0)[0]
-        edge_labels = edge_label.unbind(dim=-1)
+        edge_labels = list(edge_label.unbind(dim=-1))
         edge_labels = [one_hot(e) for e in edge_labels]
         if len(edge_labels) == 1:
             edge_label = edge_labels[0]
diff --git a/torch_geometric/loader/utils.py b/torch_geometric/loader/utils.py
index 70828d0c4..ac6643c1c 100644
--- a/torch_geometric/loader/utils.py
+++ b/torch_geometric/loader/utils.py
@@ -8,7 +8,6 @@ import torch
 from torch import Tensor
 
 import torch_geometric.typing
-from torch_geometric import EdgeIndex
 from torch_geometric.data import (
     Data,
     FeatureStore,
@@ -105,13 +104,15 @@ def filter_edge_store_(store: EdgeStorage, out_store: EdgeStorage, row: Tensor,
     # which represents the new graph as denoted by `(row, col)`:
     for key, value in store.items():
         if key == 'edge_index':
+            edge_index = torch.stack([row, col], dim=0).to(value.device)
             # TODO Integrate `EdgeIndex` into `custom_store`.
-            out_store.edge_index = EdgeIndex(
-                torch.stack([row, col], dim=0).to(value.device),
-                sparse_size=out_store.size(),
-                sort_order='col',
-                # TODO Support `is_undirected`.
-            )
+            # edge_index = EdgeIndex(
+            #     torch.stack([row, col], dim=0).to(value.device),
+            #     sparse_size=out_store.size(),
+            #     sort_order='col',
+            #     # TODO Support `is_undirected`.
+            # )
+            out_store.edge_index = edge_index
 
         elif key == 'adj_t':
             # NOTE: We expect `(row, col)` to be sorted by `col` (CSC layout).


Test patch:
diff --git a/.github/workflows/prev_testing.yml b/.github/workflows/prev_testing.yml
index 9d64df021..bb73168d0 100644
--- a/.github/workflows/prev_testing.yml
+++ b/.github/workflows/prev_testing.yml
@@ -14,7 +14,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        torch-version: [1.13.0, 2.0.0, 2.1.0]
+        torch-version: [1.13.0, 2.0.0, 2.1.0, 2.2.0]
         include:
           - torch-version: 1.13.0
             torchvision-version: 0.14.0
@@ -22,6 +22,8 @@ jobs:
             torchvision-version: 0.15.0
           - torch-version: 2.1.0
             torchvision-version: 0.16.0
+          - torch-version: 2.2.0
+            torchvision-version: 0.17.0
 
     steps:
       - name: Checkout repository
diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml
index 8e55355f2..37feabb10 100644
--- a/.github/workflows/testing.yml
+++ b/.github/workflows/testing.yml
@@ -1,4 +1,4 @@
-name: Testing PyTorch 2.2
+name: Testing PyTorch 2.3
 
 on:  # yamllint disable-line rule:truthy
   push:
diff --git a/test/loader/test_neighbor_loader.py b/test/loader/test_neighbor_loader.py
index 632e92a77..0de65855e 100644
--- a/test/loader/test_neighbor_loader.py
+++ b/test/loader/test_neighbor_loader.py
@@ -110,11 +110,12 @@ def test_homo_neighbor_loader_basic(
         assert batch.n_id.size() == (batch.num_nodes, )
         assert batch.input_id.numel() == batch.batch_size == 20
         assert batch.x.min() >= 0 and batch.x.max() < 100
-        assert isinstance(batch.edge_index, EdgeIndex)
-        batch.edge_index.validate()
-        size = (batch.num_nodes, batch.num_nodes)
-        assert batch.edge_index.sparse_size() == size
-        assert batch.edge_index.sort_order == 'col'
+        # TODO Re-enable once `EdgeIndex` is stable.
+        assert not isinstance(batch.edge_index, EdgeIndex)
+        # batch.edge_index.validate()
+        # size = (batch.num_nodes, batch.num_nodes)
+        # assert batch.edge_index.sparse_size() == size
+        # assert batch.edge_index.sort_order == 'col'
         assert batch.edge_index.device == device
         assert batch.edge_index.min() >= 0
         assert batch.edge_index.max() < batch.num_nodes
@@ -224,11 +225,12 @@ def test_hetero_neighbor_loader_basic(subgraph_type, dtype):
 
         for edge_type, edge_index in batch.edge_index_dict.items():
             src, _, dst = edge_type
-            assert isinstance(edge_index, EdgeIndex)
-            edge_index.validate()
-            size = (batch[src].num_nodes, batch[dst].num_nodes)
-            assert edge_index.sparse_size() == size
-            assert edge_index.sort_order == 'col'
+            # TODO Re-enable once `EdgeIndex` is stable.
+            assert not isinstance(edge_index, EdgeIndex)
+            # edge_index.validate()
+            # size = (batch[src].num_nodes, batch[dst].num_nodes)
+            # assert edge_index.sparse_size() == size
+            # assert edge_index.sort_order == 'col'
 
         row, col = batch['paper', 'paper'].edge_index
         assert row.min() >= 0 and row.max() < batch['paper'].num_nodes
diff --git a/test/nn/test_compile_conv.py b/test/nn/test_compile_conv.py
index 2fdf8f850..c050be408 100644
--- a/test/nn/test_compile_conv.py
+++ b/test/nn/test_compile_conv.py
@@ -52,7 +52,7 @@ def test_compile_conv(device, Conv):
 @withDevice
 @onlyLinux
 @onlyFullTest
-@withPackage('torch==2.3.0')
+@withPackage('torch==2.3')
 @pytest.mark.parametrize('Conv', [GCNConv, SAGEConv])
 def test_compile_conv_edge_index(device, Conv):
     import torch._dynamo as dynamo
diff --git a/test/test_edge_index.py b/test/test_edge_index.py
index 81d6ad578..ca045b988 100644
--- a/test/test_edge_index.py
+++ b/test/test_edge_index.py
@@ -96,6 +96,10 @@ def test_identity(dtype, device, is_undirected):
     assert out.sort_order == adj.sort_order
     assert out.is_undirected == adj.is_undirected
 
+    out = EdgeIndex(adj, sparse_size=(4, 4), sort_order='row')
+    assert out.sparse_size() == (4, 4)
+    assert out.sort_order == 'row'
+
 
 @withCUDA
 @pytest.mark.parametrize('dtype', DTYPES)
@@ -704,12 +708,14 @@ def test_add(dtype, device, is_undirected):
     assert not out.is_undirected
     assert out.sparse_size() == (6, 6)
 
-    # TODO Bring back.
-    # adj += 2
-    # assert isinstance(adj, EdgeIndex)
-    # assert adj.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
-    # assert adj.is_undirected == is_undirected
-    # assert adj.sparse_size() == (5, 5)
+    adj += 2
+    assert isinstance(adj, EdgeIndex)
+    assert adj.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
+    assert adj.is_undirected == is_undirected
+    assert adj.sparse_size() == (5, 5)
+
+    with pytest.raises(RuntimeError, match="can't be cast"):
+        adj += 2.5
 
 
 @withCUDA
@@ -749,12 +755,14 @@ def test_sub(dtype, device, is_undirected):
     assert not out.is_undirected
     assert out.sparse_size() == (None, None)
 
-    # TODO Bring back.
-    # adj -= 2
-    # assert isinstance(adj, EdgeIndex)
-    # assert adj.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
-    # assert adj.is_undirected == is_undirected
-    # assert adj.sparse_size() == (5, 5)
+    adj -= 2
+    assert isinstance(adj, EdgeIndex)
+    assert adj.equal(tensor([[2, 3, 3, 4], [3, 2, 4, 3]], device=device))
+    assert adj.is_undirected == is_undirected
+    assert adj.sparse_size() == (5, 5)
+
+    with pytest.raises(RuntimeError, match="can't be cast"):
+        adj -= 2.5
 
 
 @withCUDA
@@ -1171,7 +1179,7 @@ def test_torch_script():
 
 
 @onlyLinux
-@withPackage('torch==2.3.0')
+@withPackage('torch==2.3')
 def test_compile_basic():
     import torch._dynamo as dynamo
 
@@ -1202,7 +1210,7 @@ def test_compile_basic():
 
 
 @onlyLinux
-@withPackage('torch==2.3.0')
+@withPackage('torch==2.3')
 @pytest.mark.skip(reason="Does not work currently")
 def test_compile_create_edge_index():
     import torch._dynamo as dynamo