Repository: pyg-team/pytorch_geometric
Pull request #7754 (2023-07-17 06:45:16)

Fix `bf16` data type failure in `spmm`
This PR is to fix the data type failure in bf16 with sparse graph. It will convert src to other's dtype in spmm.py.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7754?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7754](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7754?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (13c8762) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/6995c54eb22195bc1771df2ac6c7e639e27c1c6a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (6995c54) will **decrease** coverage by `0.36%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7754      +/-   ##
==========================================
- Coverage   91.93%   91.58%   -0.36%     
==========================================
  Files         452      452              
  Lines       25537    25534       -3     
==========================================
- Hits        23478    23385      -93     
- Misses       2059     2149      +90     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7754?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/spmm.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7754?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NwbW0ucHk=) | `56.25% <100.00%> (-34.38%)` | :arrow_down: |

... and [19 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7754/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/.gitattributes b/.gitattributes
new file mode 100644
index 000000000..d9bd16b09
--- /dev/null
+++ b/.gitattributes
@@ -0,0 +1,1 @@
+*.py text eol=lf
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e9eaa0d29..9bcc1c405 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 possibility to run inference benchmarks on XPU device ([#7705](https://github.com/pyg-team/pytorch_geometric/pull/7705))
 - Added `HeteroData` support in `to_networkx` ([#7713](https://github.com/pyg-team/pytorch_geometric/pull/7713))
 - Added `FlopsCount` support via `fvcore` ([#7693](https://github.com/pyg-team/pytorch_geometric/pull/7693))
 - Added back support for PyTorch >= 1.11.0 ([#7656](https://github.com/pyg-team/pytorch_geometric/pull/7656))
diff --git a/benchmark/inference/inference_benchmark.py b/benchmark/inference/inference_benchmark.py
index 380c1d53f..22d6dd7bf 100644
--- a/benchmark/inference/inference_benchmark.py
+++ b/benchmark/inference/inference_benchmark.py
@@ -16,7 +16,12 @@ from benchmark.utils import (
 )
 from torch_geometric.loader import NeighborLoader
 from torch_geometric.nn import PNAConv
-from torch_geometric.profile import rename_profile_file, timeit, torch_profile
+from torch_geometric.profile import (
+    rename_profile_file,
+    timeit,
+    torch_profile,
+    xpu_profile,
+)
 
 supported_sets = {
     'ogbn-mag': ['rgat', 'rgcn'],
@@ -42,11 +47,23 @@ def run(args: argparse.ArgumentParser):
         warnings.warn("Cannot write profile data to CSV because profiling is "
                       "disabled")
 
-    # cuda device is not suitable for full batch mode
-    device = torch.device(
-        'cuda' if not args.full_batch and torch.cuda.is_available() else 'cpu')
+    if args.device == 'xpu':
+        try:
+            import intel_extension_for_pytorch as ipex
+        except ImportError:
+            raise RuntimeError('XPU device requires IPEX to be installed')
+
+    if ((args.device == 'cuda' and not torch.cuda.is_available())
+            or (args.device == 'xpu' and not torch.xpu.is_available())):
+        raise RuntimeError(f'{args.device.upper()} is not available')
+
+    if args.device == 'cuda' and args.full_batch:
+        raise RuntimeError('CUDA device is not suitable for full batch mode')
+
+    device = torch.device(args.device)
 
     print('BENCHMARK STARTS')
+    print(f'Running on {args.device.upper()}')
     for dataset_name in args.datasets:
         assert dataset_name in supported_sets.keys(
         ), f"Dataset {dataset_name} isn't supported."
@@ -66,11 +83,17 @@ def run(args: argparse.ArgumentParser):
         if args.num_layers != [1] and not hetero and args.num_steps != -1:
             raise ValueError("Layer-wise inference requires `steps=-1`")
 
-        if torch.cuda.is_available():
+        if args.device == 'cuda':
             amp = torch.cuda.amp.autocast(enabled=False)
+        elif args.device == 'xpu':
+            amp = torch.xpu.amp.autocast(enabled=False)
         else:
             amp = torch.cpu.amp.autocast(enabled=args.bf16)
 
+        if args.device == 'xpu' and args.warmup < 1:
+            print('XPU device requires warmup - setting warmup=1')
+            args.warmup = 1
+
         inputs_channels = data[
             'paper'].num_features if dataset_name == 'ogbn-mag' \
             else dataset.num_features
@@ -163,6 +186,8 @@ def run(args: argparse.ArgumentParser):
                             state_dict = torch.load(args.ckpt_path)
                             model.load_state_dict(state_dict)
                         model.eval()
+                        if args.device == 'xpu':
+                            model = ipex.optimize(model)
 
                         # Define context manager parameters:
                         if args.cpu_affinity and with_loader:
@@ -170,9 +195,13 @@ def run(args: argparse.ArgumentParser):
                                 args.loader_cores)
                         else:
                             cpu_affinity = nullcontext()
-                        profile = torch_profile(
-                            args.export_chrome_trace, csv_data,
-                            args.write_csv) if args.profile else nullcontext()
+                        if args.profile and args.device == 'xpu':
+                            profile = xpu_profile(args.export_chrome_trace)
+                        elif args.profile:
+                            profile = torch_profile(args.export_chrome_trace,
+                                                    csv_data, args.write_csv)
+                        else:
+                            profile = nullcontext()
                         itt = emit_itt(
                         ) if args.vtune_profile else nullcontext()
 
@@ -256,6 +285,8 @@ if __name__ == '__main__':
     argparser = argparse.ArgumentParser('GNN inference benchmark')
     add = argparser.add_argument
 
+    add('--device', choices=['cpu', 'cuda', 'xpu'], default='cpu',
+        help='Device to run benchmark on')
     add('--datasets', nargs='+',
         default=['ogbn-mag', 'ogbn-products', 'Reddit'], type=str)
     add('--use-sparse-tensor', action='store_true',
diff --git a/examples/pytorch_lightning/relational_gnn.py b/examples/pytorch_lightning/relational_gnn.py
index 042ebdf60..fa7b264c7 100644
--- a/examples/pytorch_lightning/relational_gnn.py
+++ b/examples/pytorch_lightning/relational_gnn.py
@@ -21,7 +21,7 @@ class GNN(torch.nn.Module):
     def __init__(self, hidden_channels: int, out_channels: int,
                  dropout: float):
         super().__init__()
-        self.dropout = dropout
+        self.dropout = torch.nn.Dropout(p=dropout)
 
         self.conv1 = SAGEConv((-1, -1), hidden_channels)
         self.conv2 = SAGEConv((-1, -1), hidden_channels)
@@ -29,9 +29,9 @@ class GNN(torch.nn.Module):
 
     def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
         x = self.conv1(x, edge_index).relu()
-        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.dropout(x)
         x = self.conv2(x, edge_index).relu()
-        x = F.dropout(x, p=self.dropout, training=self.training)
+        x = self.dropout(x)
         return self.lin(x)
 
 
diff --git a/examples/rect.py b/examples/rect.py
index 6d59991c6..a89e6ff1f 100644
--- a/examples/rect.py
+++ b/examples/rect.py
@@ -1,66 +1,66 @@
-import argparse
-import copy
-import os.path as osp
-
-import torch
-from sklearn.linear_model import LogisticRegression
-
-import torch_geometric.transforms as T
-from torch_geometric.datasets import Planetoid
-from torch_geometric.nn import RECT_L
-
-# RECT focuses on the zero-shot, i.e. completely-imbalanced label setting:
-# For this, we first remove "unseen" classes from the training set and train a
-# RECT (or more specifically its supervised part RECT-L) model in the zero-shot
-# label scenario. Lastly, we train a simple classifier to evaluate the final
-# performance of the embeddings based on the original labels.
-
-# Datasets              Citeseer             Cora          Pubmed
-# Unseen Classes  [1, 2, 5]  [3, 4]  [1, 2, 3]  [3, 4, 6]  [2]
-# RECT-L          66.30      68.20   74.60      71.20      75.30
-# GCN             51.80      55.70   55.80      57.10      59.80
-# NodeFeats       61.40      61.40   57.50      57.50      73.10
-
-parser = argparse.ArgumentParser()
-parser.add_argument('--dataset', type=str, default='Cora',
-                    choices=['Cora', 'CiteSeer', 'PubMed'])
-parser.add_argument('--unseen-classes', type=int, nargs='*', default=[1, 2, 3])
-args = parser.parse_args()
-
-path = osp.join(osp.dirname(osp.realpath(__file__)), '../data/Planetoid')
-train_mask_original = Planetoid(path, args.dataset)[0].train_mask.clone()
-transform = T.Compose([
-    T.NormalizeFeatures(),
-    T.SVDFeatureReduction(200),
-    T.GDC(),
-])
-dataset = Planetoid(path, args.dataset, transform=transform)
-data = dataset[0]
-zs_data = T.RemoveTrainingClasses(args.unseen_classes)(copy.copy(data))
-
-model = RECT_L(200, 200, normalize=False, dropout=0.0)
-zs_data.y = model.get_semantic_labels(zs_data.x, zs_data.y, zs_data.train_mask)
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-model, zs_data = model.to(device), zs_data.to(device)
-
-criterion = torch.nn.MSELoss(reduction='sum')
-optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4)
-
-model.train()
-for epoch in range(1, 201):
-    optimizer.zero_grad()
-    out = model(zs_data.x, zs_data.edge_index, zs_data.edge_attr)
-    loss = criterion(out[zs_data.train_mask], zs_data.y)
-    loss.backward()
-    optimizer.step()
-    print(f'Epoch {epoch:03d}, Loss {loss:.4f}')
-
-model.eval()
-with torch.no_grad():
-    h = model.embed(zs_data.x, zs_data.edge_index, zs_data.edge_attr).cpu()
-
-reg = LogisticRegression()
-reg.fit(h[data.train_mask].numpy(), data.y[data.train_mask].numpy())
-test_acc = reg.score(h[data.test_mask].numpy(), data.y[data.test_mask].numpy())
-print(f'Test Acc: {test_acc:.4f}')
+import argparse
+import copy
+import os.path as osp
+
+import torch
+from sklearn.linear_model import LogisticRegression
+
+import torch_geometric.transforms as T
+from torch_geometric.datasets import Planetoid
+from torch_geometric.nn import RECT_L
+
+# RECT focuses on the zero-shot, i.e. completely-imbalanced label setting:
+# For this, we first remove "unseen" classes from the training set and train a
+# RECT (or more specifically its supervised part RECT-L) model in the zero-shot
+# label scenario. Lastly, we train a simple classifier to evaluate the final
+# performance of the embeddings based on the original labels.
+
+# Datasets              Citeseer             Cora          Pubmed
+# Unseen Classes  [1, 2, 5]  [3, 4]  [1, 2, 3]  [3, 4, 6]  [2]
+# RECT-L          66.30      68.20   74.60      71.20      75.30
+# GCN             51.80      55.70   55.80      57.10      59.80
+# NodeFeats       61.40      61.40   57.50      57.50      73.10
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dataset', type=str, default='Cora',
+                    choices=['Cora', 'CiteSeer', 'PubMed'])
+parser.add_argument('--unseen-classes', type=int, nargs='*', default=[1, 2, 3])
+args = parser.parse_args()
+
+path = osp.join(osp.dirname(osp.realpath(__file__)), '../data/Planetoid')
+train_mask_original = Planetoid(path, args.dataset)[0].train_mask.clone()
+transform = T.Compose([
+    T.NormalizeFeatures(),
+    T.SVDFeatureReduction(200),
+    T.GDC(),
+])
+dataset = Planetoid(path, args.dataset, transform=transform)
+data = dataset[0]
+zs_data = T.RemoveTrainingClasses(args.unseen_classes)(copy.copy(data))
+
+model = RECT_L(200, 200, normalize=False, dropout=0.0)
+zs_data.y = model.get_semantic_labels(zs_data.x, zs_data.y, zs_data.train_mask)
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+model, zs_data = model.to(device), zs_data.to(device)
+
+criterion = torch.nn.MSELoss(reduction='sum')
+optimizer = torch.optim.Adam(model.parameters(), lr=0.001, weight_decay=5e-4)
+
+model.train()
+for epoch in range(1, 201):
+    optimizer.zero_grad()
+    out = model(zs_data.x, zs_data.edge_index, zs_data.edge_attr)
+    loss = criterion(out[zs_data.train_mask], zs_data.y)
+    loss.backward()
+    optimizer.step()
+    print(f'Epoch {epoch:03d}, Loss {loss:.4f}')
+
+model.eval()
+with torch.no_grad():
+    h = model.embed(zs_data.x, zs_data.edge_index, zs_data.edge_attr).cpu()
+
+reg = LogisticRegression()
+reg.fit(h[data.train_mask].numpy(), data.y[data.train_mask].numpy())
+test_acc = reg.score(h[data.test_mask].numpy(), data.y[data.test_mask].numpy())
+print(f'Test Acc: {test_acc:.4f}')
diff --git a/examples/rgcn_link_pred.py b/examples/rgcn_link_pred.py
index 71c768bfb..fd47d254c 100644
--- a/examples/rgcn_link_pred.py
+++ b/examples/rgcn_link_pred.py
@@ -16,9 +16,11 @@ from tqdm import tqdm
 from torch_geometric.datasets import RelLinkPredDataset
 from torch_geometric.nn import GAE, RGCNConv
 
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'RLPD')
 dataset = RelLinkPredDataset(path, 'FB15k-237')
-data = dataset[0]
+data = dataset[0].to(device)
 
 
 class RGCNEncoder(torch.nn.Module):
@@ -60,10 +62,10 @@ class DistMultDecoder(torch.nn.Module):
 
 
 model = GAE(
-    RGCNEncoder(data.num_nodes, hidden_channels=500,
-                num_relations=dataset.num_relations),
-    DistMultDecoder(dataset.num_relations // 2, hidden_channels=500),
-)
+    RGCNEncoder(data.num_nodes, 500, dataset.num_relations),
+    DistMultDecoder(dataset.num_relations // 2, 500),
+).to(device)
+
 optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
 
 
@@ -73,8 +75,10 @@ def negative_sampling(edge_index, num_nodes):
     mask_2 = ~mask_1
 
     neg_edge_index = edge_index.clone()
-    neg_edge_index[0, mask_1] = torch.randint(num_nodes, (mask_1.sum(), ))
-    neg_edge_index[1, mask_2] = torch.randint(num_nodes, (mask_2.sum(), ))
+    neg_edge_index[0, mask_1] = torch.randint(num_nodes, (mask_1.sum(), ),
+                                              device=neg_edge_index.device)
+    neg_edge_index[1, mask_2] = torch.randint(num_nodes, (mask_2.sum(), ),
+                                              device=neg_edge_index.device)
     return neg_edge_index
 
 
diff --git a/torch_geometric/datasets/movie_lens.py b/torch_geometric/datasets/movie_lens.py
index bc872625b..9316a7849 100644
--- a/torch_geometric/datasets/movie_lens.py
+++ b/torch_geometric/datasets/movie_lens.py
@@ -1,99 +1,99 @@
-import os
-import os.path as osp
-from typing import Callable, List, Optional
-
-import torch
-
-from torch_geometric.data import (
-    HeteroData,
-    InMemoryDataset,
-    download_url,
-    extract_zip,
-)
-
-
-class MovieLens(InMemoryDataset):
-    r"""A heterogeneous rating dataset, assembled by GroupLens Research from
-    the `MovieLens web site <https://movielens.org>`_, consisting of nodes of
-    type :obj:`"movie"` and :obj:`"user"`.
-    User ratings for movies are available as ground truth labels for the edges
-    between the users and the movies :obj:`("user", "rates", "movie")`.
-
-    Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.HeteroData` object and returns a
-            transformed version. The data object will be transformed before
-            every access. (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.HeteroData` object and returns a
-            transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-        model_name (str): Name of model used to transform movie titles to node
-            features. The model comes from the`Huggingface SentenceTransformer
-            <https://huggingface.co/sentence-transformers>`_.
-    """
-    url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
-
-    def __init__(
-        self,
-        root: str,
-        transform: Optional[Callable] = None,
-        pre_transform: Optional[Callable] = None,
-        model_name: Optional[str] = 'all-MiniLM-L6-v2',
-    ):
-        self.model_name = model_name
-        super().__init__(root, transform, pre_transform)
-        self.data, self.slices = torch.load(self.processed_paths[0])
-
-    @property
-    def raw_file_names(self) -> List[str]:
-        return [
-            osp.join('ml-latest-small', 'movies.csv'),
-            osp.join('ml-latest-small', 'ratings.csv'),
-        ]
-
-    @property
-    def processed_file_names(self) -> str:
-        return f'data_{self.model_name}.pt'
-
-    def download(self):
-        path = download_url(self.url, self.raw_dir)
-        extract_zip(path, self.raw_dir)
-        os.remove(path)
-
-    def process(self):
-        import pandas as pd
-        from sentence_transformers import SentenceTransformer
-
-        data = HeteroData()
-
-        df = pd.read_csv(self.raw_paths[0], index_col='movieId')
-        movie_mapping = {idx: i for i, idx in enumerate(df.index)}
-
-        genres = df['genres'].str.get_dummies('|').values
-        genres = torch.from_numpy(genres).to(torch.float)
-
-        model = SentenceTransformer(self.model_name)
-        with torch.no_grad():
-            emb = model.encode(df['title'].values, show_progress_bar=True,
-                               convert_to_tensor=True).cpu()
-
-        data['movie'].x = torch.cat([emb, genres], dim=-1)
-
-        df = pd.read_csv(self.raw_paths[1])
-        user_mapping = {idx: i for i, idx in enumerate(df['userId'].unique())}
-        data['user'].num_nodes = len(user_mapping)
-
-        src = [user_mapping[idx] for idx in df['userId']]
-        dst = [movie_mapping[idx] for idx in df['movieId']]
-        edge_index = torch.tensor([src, dst])
-
-        rating = torch.from_numpy(df['rating'].values).to(torch.long)
-        data['user', 'rates', 'movie'].edge_index = edge_index
-        data['user', 'rates', 'movie'].edge_label = rating
-
-        if self.pre_transform is not None:
-            data = self.pre_transform(data)
-
-        torch.save(self.collate([data]), self.processed_paths[0])
+import os
+import os.path as osp
+from typing import Callable, List, Optional
+
+import torch
+
+from torch_geometric.data import (
+    HeteroData,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+
+
+class MovieLens(InMemoryDataset):
+    r"""A heterogeneous rating dataset, assembled by GroupLens Research from
+    the `MovieLens web site <https://movielens.org>`_, consisting of nodes of
+    type :obj:`"movie"` and :obj:`"user"`.
+    User ratings for movies are available as ground truth labels for the edges
+    between the users and the movies :obj:`("user", "rates", "movie")`.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            every access. (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+        model_name (str): Name of model used to transform movie titles to node
+            features. The model comes from the`Huggingface SentenceTransformer
+            <https://huggingface.co/sentence-transformers>`_.
+    """
+    url = 'https://files.grouplens.org/datasets/movielens/ml-latest-small.zip'
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+        model_name: Optional[str] = 'all-MiniLM-L6-v2',
+    ):
+        self.model_name = model_name
+        super().__init__(root, transform, pre_transform)
+        self.data, self.slices = torch.load(self.processed_paths[0])
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return [
+            osp.join('ml-latest-small', 'movies.csv'),
+            osp.join('ml-latest-small', 'ratings.csv'),
+        ]
+
+    @property
+    def processed_file_names(self) -> str:
+        return f'data_{self.model_name}.pt'
+
+    def download(self):
+        path = download_url(self.url, self.raw_dir)
+        extract_zip(path, self.raw_dir)
+        os.remove(path)
+
+    def process(self):
+        import pandas as pd
+        from sentence_transformers import SentenceTransformer
+
+        data = HeteroData()
+
+        df = pd.read_csv(self.raw_paths[0], index_col='movieId')
+        movie_mapping = {idx: i for i, idx in enumerate(df.index)}
+
+        genres = df['genres'].str.get_dummies('|').values
+        genres = torch.from_numpy(genres).to(torch.float)
+
+        model = SentenceTransformer(self.model_name)
+        with torch.no_grad():
+            emb = model.encode(df['title'].values, show_progress_bar=True,
+                               convert_to_tensor=True).cpu()
+
+        data['movie'].x = torch.cat([emb, genres], dim=-1)
+
+        df = pd.read_csv(self.raw_paths[1])
+        user_mapping = {idx: i for i, idx in enumerate(df['userId'].unique())}
+        data['user'].num_nodes = len(user_mapping)
+
+        src = [user_mapping[idx] for idx in df['userId']]
+        dst = [movie_mapping[idx] for idx in df['movieId']]
+        edge_index = torch.tensor([src, dst])
+
+        rating = torch.from_numpy(df['rating'].values).to(torch.long)
+        data['user', 'rates', 'movie'].edge_index = edge_index
+        data['user', 'rates', 'movie'].edge_label = rating
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        torch.save(self.collate([data]), self.processed_paths[0])
diff --git a/torch_geometric/datasets/movie_lens_100k.py b/torch_geometric/datasets/movie_lens_100k.py
index b7b8c6001..cf83dca51 100644
--- a/torch_geometric/datasets/movie_lens_100k.py
+++ b/torch_geometric/datasets/movie_lens_100k.py
@@ -1,179 +1,179 @@
-import os
-import os.path as osp
-import shutil
-from typing import Callable, List, Optional
-
-import torch
-
-from torch_geometric.data import (
-    HeteroData,
-    InMemoryDataset,
-    download_url,
-    extract_zip,
-)
-
-MOVIE_HEADERS = [
-    "movieId", "title", "releaseDate", "videoReleaseDate", "IMDb URL",
-    "unknown", "Action", "Adventure", "Animation", "Children's", "Comedy",
-    "Crime", "Documentary", "Drama", "Fantasy", "Film-Noir", "Horror",
-    "Musical", "Mystery", "Romance", "Sci-Fi", "Thriller", "War", "Western"
-]
-USER_HEADERS = ["userId", "age", "gender", "occupation", "zipCode"]
-RATING_HEADERS = ["userId", "movieId", "rating", "timestamp"]
-
-
-class MovieLens100K(InMemoryDataset):
-    r"""The MovieLens 100K heterogeneous rating dataset, assembled by GroupLens
-    Research from the `MovieLens web site <https://movielens.org>`__,
-    consisting of movies (1,682 nodes) and users (943 nodes) with 100K
-    ratings between them.
-    User ratings for movies are available as ground truth labels.
-    Features of users and movies are encoded according to the `"Inductive
-    Matrix Completion Based on Graph Neural Networks"
-    <https://arxiv.org/abs/1904.12058>`__ paper.
-
-    Args:
-        root (str): Root directory where the dataset should be saved.
-        transform (callable, optional): A function/transform that takes in an
-            :obj:`torch_geometric.data.HeteroData` object and returns a
-            transformed version. The data object will be transformed before
-            every access. (default: :obj:`None`)
-        pre_transform (callable, optional): A function/transform that takes in
-            an :obj:`torch_geometric.data.HeteroData` object and returns a
-            transformed version. The data object will be transformed before
-            being saved to disk. (default: :obj:`None`)
-
-    **STATS:**
-
-    .. list-table::
-        :widths: 20 10 10 10
-        :header-rows: 1
-
-        * - Node/Edge Type
-          - #nodes/#edges
-          - #features
-          - #tasks
-        * - Movie
-          - 1,682
-          - 18
-          -
-        * - User
-          - 943
-          - 24
-          -
-        * - User-Movie
-          - 80,000
-          - 1
-          - 1
-    """
-    url = 'https://files.grouplens.org/datasets/movielens/ml-100k.zip'
-
-    def __init__(
-        self,
-        root: str,
-        transform: Optional[Callable] = None,
-        pre_transform: Optional[Callable] = None,
-    ):
-        super().__init__(root, transform, pre_transform)
-        self.load(self.processed_paths[0], data_cls=HeteroData)
-
-    @property
-    def raw_file_names(self) -> List[str]:
-        return ['u.item', 'u.user', 'u1.base', 'u1.test']
-
-    @property
-    def processed_file_names(self) -> str:
-        return 'data.pt'
-
-    def download(self):
-        path = download_url(self.url, self.root)
-        extract_zip(path, self.root)
-        os.remove(path)
-        folder = osp.join(self.root, 'ml-100k')
-        shutil.rmtree(self.raw_dir)
-        os.rename(folder, self.raw_dir)
-
-    def process(self):
-        import pandas as pd
-
-        data = HeteroData()
-
-        # Process movie data:
-        df = pd.read_csv(
-            self.raw_paths[0],
-            sep='|',
-            header=None,
-            names=MOVIE_HEADERS,
-            index_col='movieId',
-            encoding='ISO-8859-1',
-        )
-        movie_mapping = {idx: i for i, idx in enumerate(df.index)}
-
-        x = df[MOVIE_HEADERS[6:]].values
-        data['movie'].x = torch.from_numpy(x).to(torch.float)
-
-        # Process user data:
-        df = pd.read_csv(
-            self.raw_paths[1],
-            sep='|',
-            header=None,
-            names=USER_HEADERS,
-            index_col='userId',
-            encoding='ISO-8859-1',
-        )
-        user_mapping = {idx: i for i, idx in enumerate(df.index)}
-
-        age = df['age'].values / df['age'].values.max()
-        age = torch.from_numpy(age).to(torch.float).view(-1, 1)
-
-        gender = df['gender'].str.get_dummies().values
-        gender = torch.from_numpy(gender).to(torch.float)
-
-        occupation = df['occupation'].str.get_dummies().values
-        occupation = torch.from_numpy(occupation).to(torch.float)
-
-        data['user'].x = torch.cat([age, gender, occupation], dim=-1)
-
-        # Process rating data for training:
-        df = pd.read_csv(
-            self.raw_paths[2],
-            sep='\t',
-            header=None,
-            names=RATING_HEADERS,
-        )
-
-        src = [user_mapping[idx] for idx in df['userId']]
-        dst = [movie_mapping[idx] for idx in df['movieId']]
-        edge_index = torch.tensor([src, dst])
-        data['user', 'rates', 'movie'].edge_index = edge_index
-
-        rating = torch.from_numpy(df['rating'].values).to(torch.long)
-        data['user', 'rates', 'movie'].rating = rating
-
-        time = torch.from_numpy(df['timestamp'].values)
-        data['user', 'rates', 'movie'].time = time
-
-        data['movie', 'rated_by', 'user'].edge_index = edge_index.flip([0])
-        data['movie', 'rated_by', 'user'].rating = rating
-        data['movie', 'rated_by', 'user'].time = time
-
-        # Process rating data for testing:
-        df = pd.read_csv(
-            self.raw_paths[3],
-            sep='\t',
-            header=None,
-            names=RATING_HEADERS,
-        )
-
-        src = [user_mapping[idx] for idx in df['userId']]
-        dst = [movie_mapping[idx] for idx in df['movieId']]
-        edge_label_index = torch.tensor([src, dst])
-        data['user', 'rates', 'movie'].edge_label_index = edge_label_index
-
-        edge_label = torch.from_numpy(df['rating'].values).to(torch.float)
-        data['user', 'rates', 'movie'].edge_label = edge_label
-
-        if self.pre_transform is not None:
-            data = self.pre_transform(data)
-
-        self.save([data], self.processed_paths[0])
+import os
+import os.path as osp
+import shutil
+from typing import Callable, List, Optional
+
+import torch
+
+from torch_geometric.data import (
+    HeteroData,
+    InMemoryDataset,
+    download_url,
+    extract_zip,
+)
+
+MOVIE_HEADERS = [
+    "movieId", "title", "releaseDate", "videoReleaseDate", "IMDb URL",
+    "unknown", "Action", "Adventure", "Animation", "Children's", "Comedy",
+    "Crime", "Documentary", "Drama", "Fantasy", "Film-Noir", "Horror",
+    "Musical", "Mystery", "Romance", "Sci-Fi", "Thriller", "War", "Western"
+]
+USER_HEADERS = ["userId", "age", "gender", "occupation", "zipCode"]
+RATING_HEADERS = ["userId", "movieId", "rating", "timestamp"]
+
+
+class MovieLens100K(InMemoryDataset):
+    r"""The MovieLens 100K heterogeneous rating dataset, assembled by GroupLens
+    Research from the `MovieLens web site <https://movielens.org>`__,
+    consisting of movies (1,682 nodes) and users (943 nodes) with 100K
+    ratings between them.
+    User ratings for movies are available as ground truth labels.
+    Features of users and movies are encoded according to the `"Inductive
+    Matrix Completion Based on Graph Neural Networks"
+    <https://arxiv.org/abs/1904.12058>`__ paper.
+
+    Args:
+        root (str): Root directory where the dataset should be saved.
+        transform (callable, optional): A function/transform that takes in an
+            :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            every access. (default: :obj:`None`)
+        pre_transform (callable, optional): A function/transform that takes in
+            an :obj:`torch_geometric.data.HeteroData` object and returns a
+            transformed version. The data object will be transformed before
+            being saved to disk. (default: :obj:`None`)
+
+    **STATS:**
+
+    .. list-table::
+        :widths: 20 10 10 10
+        :header-rows: 1
+
+        * - Node/Edge Type
+          - #nodes/#edges
+          - #features
+          - #tasks
+        * - Movie
+          - 1,682
+          - 18
+          -
+        * - User
+          - 943
+          - 24
+          -
+        * - User-Movie
+          - 80,000
+          - 1
+          - 1
+    """
+    url = 'https://files.grouplens.org/datasets/movielens/ml-100k.zip'
+
+    def __init__(
+        self,
+        root: str,
+        transform: Optional[Callable] = None,
+        pre_transform: Optional[Callable] = None,
+    ):
+        super().__init__(root, transform, pre_transform)
+        self.load(self.processed_paths[0], data_cls=HeteroData)
+
+    @property
+    def raw_file_names(self) -> List[str]:
+        return ['u.item', 'u.user', 'u1.base', 'u1.test']
+
+    @property
+    def processed_file_names(self) -> str:
+        return 'data.pt'
+
+    def download(self):
+        path = download_url(self.url, self.root)
+        extract_zip(path, self.root)
+        os.remove(path)
+        folder = osp.join(self.root, 'ml-100k')
+        shutil.rmtree(self.raw_dir)
+        os.rename(folder, self.raw_dir)
+
+    def process(self):
+        import pandas as pd
+
+        data = HeteroData()
+
+        # Process movie data:
+        df = pd.read_csv(
+            self.raw_paths[0],
+            sep='|',
+            header=None,
+            names=MOVIE_HEADERS,
+            index_col='movieId',
+            encoding='ISO-8859-1',
+        )
+        movie_mapping = {idx: i for i, idx in enumerate(df.index)}
+
+        x = df[MOVIE_HEADERS[6:]].values
+        data['movie'].x = torch.from_numpy(x).to(torch.float)
+
+        # Process user data:
+        df = pd.read_csv(
+            self.raw_paths[1],
+            sep='|',
+            header=None,
+            names=USER_HEADERS,
+            index_col='userId',
+            encoding='ISO-8859-1',
+        )
+        user_mapping = {idx: i for i, idx in enumerate(df.index)}
+
+        age = df['age'].values / df['age'].values.max()
+        age = torch.from_numpy(age).to(torch.float).view(-1, 1)
+
+        gender = df['gender'].str.get_dummies().values
+        gender = torch.from_numpy(gender).to(torch.float)
+
+        occupation = df['occupation'].str.get_dummies().values
+        occupation = torch.from_numpy(occupation).to(torch.float)
+
+        data['user'].x = torch.cat([age, gender, occupation], dim=-1)
+
+        # Process rating data for training:
+        df = pd.read_csv(
+            self.raw_paths[2],
+            sep='\t',
+            header=None,
+            names=RATING_HEADERS,
+        )
+
+        src = [user_mapping[idx] for idx in df['userId']]
+        dst = [movie_mapping[idx] for idx in df['movieId']]
+        edge_index = torch.tensor([src, dst])
+        data['user', 'rates', 'movie'].edge_index = edge_index
+
+        rating = torch.from_numpy(df['rating'].values).to(torch.long)
+        data['user', 'rates', 'movie'].rating = rating
+
+        time = torch.from_numpy(df['timestamp'].values)
+        data['user', 'rates', 'movie'].time = time
+
+        data['movie', 'rated_by', 'user'].edge_index = edge_index.flip([0])
+        data['movie', 'rated_by', 'user'].rating = rating
+        data['movie', 'rated_by', 'user'].time = time
+
+        # Process rating data for testing:
+        df = pd.read_csv(
+            self.raw_paths[3],
+            sep='\t',
+            header=None,
+            names=RATING_HEADERS,
+        )
+
+        src = [user_mapping[idx] for idx in df['userId']]
+        dst = [movie_mapping[idx] for idx in df['movieId']]
+        edge_label_index = torch.tensor([src, dst])
+        data['user', 'rates', 'movie'].edge_label_index = edge_label_index
+
+        edge_label = torch.from_numpy(df['rating'].values).to(torch.float)
+        data['user', 'rates', 'movie'].edge_label = edge_label
+
+        if self.pre_transform is not None:
+            data = self.pre_transform(data)
+
+        self.save([data], self.processed_paths[0])
diff --git a/torch_geometric/nn/conv/han_conv.py b/torch_geometric/nn/conv/han_conv.py
index e7ba600eb..c88bd2e12 100644
--- a/torch_geometric/nn/conv/han_conv.py
+++ b/torch_geometric/nn/conv/han_conv.py
@@ -1,183 +1,183 @@
-from typing import Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor, nn
-
-from torch_geometric.nn.conv import MessagePassing
-from torch_geometric.nn.dense import Linear
-from torch_geometric.nn.inits import glorot, reset
-from torch_geometric.typing import Adj, EdgeType, Metadata, NodeType, OptTensor
-from torch_geometric.utils import softmax
-
-
-def group(
-    xs: List[Tensor],
-    q: nn.Parameter,
-    k_lin: nn.Module,
-) -> Tuple[OptTensor, OptTensor]:
-
-    if len(xs) == 0:
-        return None, None
-    else:
-        num_edge_types = len(xs)
-        out = torch.stack(xs)
-        if out.numel() == 0:
-            return out.view(0, out.size(-1)), None
-        attn_score = (q * torch.tanh(k_lin(out)).mean(1)).sum(-1)
-        attn = F.softmax(attn_score, dim=0)
-        out = torch.sum(attn.view(num_edge_types, 1, -1) * out, dim=0)
-        return out, attn
-
-
-class HANConv(MessagePassing):
-    r"""
-    The Heterogenous Graph Attention Operator from the
-    `"Heterogenous Graph Attention Network"
-    <https://arxiv.org/pdf/1903.07293.pdf>`_ paper.
-
-    .. note::
-
-        For an example of using HANConv, see `examples/hetero/han_imdb.py
-        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
-        hetero/han_imdb.py>`_.
-
-    Args:
-        in_channels (int or Dict[str, int]): Size of each input sample of every
-            node type, or :obj:`-1` to derive the size from the first input(s)
-            to the forward method.
-        out_channels (int): Size of each output sample.
-        metadata (Tuple[List[str], List[Tuple[str, str, str]]]): The metadata
-            of the heterogeneous graph, *i.e.* its node and edge types given
-            by a list of strings and a list of string triplets, respectively.
-            See :meth:`torch_geometric.data.HeteroData.metadata` for more
-            information.
-        heads (int, optional): Number of multi-head-attentions.
-            (default: :obj:`1`)
-        negative_slope (float, optional): LeakyReLU angle of the negative
-            slope. (default: :obj:`0.2`)
-        dropout (float, optional): Dropout probability of the normalized
-            attention coefficients which exposes each node to a stochastically
-            sampled neighborhood during training. (default: :obj:`0`)
-        **kwargs (optional): Additional arguments of
-            :class:`torch_geometric.nn.conv.MessagePassing`.
-    """
-    def __init__(
-        self,
-        in_channels: Union[int, Dict[str, int]],
-        out_channels: int,
-        metadata: Metadata,
-        heads: int = 1,
-        negative_slope=0.2,
-        dropout: float = 0.0,
-        **kwargs,
-    ):
-        super().__init__(aggr='add', node_dim=0, **kwargs)
-
-        if not isinstance(in_channels, dict):
-            in_channels = {node_type: in_channels for node_type in metadata[0]}
-
-        self.heads = heads
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-        self.negative_slope = negative_slope
-        self.metadata = metadata
-        self.dropout = dropout
-        self.k_lin = nn.Linear(out_channels, out_channels)
-        self.q = nn.Parameter(torch.empty(1, out_channels))
-
-        self.proj = nn.ModuleDict()
-        for node_type, in_channels in self.in_channels.items():
-            self.proj[node_type] = Linear(in_channels, out_channels)
-
-        self.lin_src = nn.ParameterDict()
-        self.lin_dst = nn.ParameterDict()
-        dim = out_channels // heads
-        for edge_type in metadata[1]:
-            edge_type = '__'.join(edge_type)
-            self.lin_src[edge_type] = nn.Parameter(torch.empty(1, heads, dim))
-            self.lin_dst[edge_type] = nn.Parameter(torch.empty(1, heads, dim))
-
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        super().reset_parameters()
-        reset(self.proj)
-        glorot(self.lin_src)
-        glorot(self.lin_dst)
-        self.k_lin.reset_parameters()
-        glorot(self.q)
-
-    def forward(
-        self,
-        x_dict: Dict[NodeType, Tensor],
-        edge_index_dict: Dict[EdgeType, Adj],
-        return_semantic_attention_weights: bool = False,
-    ) -> Union[Dict[NodeType, OptTensor], Tuple[Dict[NodeType, OptTensor],
-                                                Dict[NodeType, OptTensor]]]:
-        r"""Runs the forward pass of the module.
-
-        Args:
-            x_dict (Dict[str, torch.Tensor]): A dictionary holding node feature
-                information for each individual node type.
-            edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A
-                dictionary holding graph connectivity information for each
-                individual edge type, either as a :class:`torch.Tensor` of
-                shape :obj:`[2, num_edges]` or a
-                :class:`torch_sparse.SparseTensor`.
-            return_semantic_attention_weights (bool, optional): If set to
-                :obj:`True`, will additionally return the semantic-level
-                attention weights for each destination node type.
-                (default: :obj:`False`)
-        """
-        H, D = self.heads, self.out_channels // self.heads
-        x_node_dict, out_dict = {}, {}
-
-        # Iterate over node types:
-        for node_type, x in x_dict.items():
-            x_node_dict[node_type] = self.proj[node_type](x).view(-1, H, D)
-            out_dict[node_type] = []
-
-        # Iterate over edge types:
-        for edge_type, edge_index in edge_index_dict.items():
-            src_type, _, dst_type = edge_type
-            edge_type = '__'.join(edge_type)
-            lin_src = self.lin_src[edge_type]
-            lin_dst = self.lin_dst[edge_type]
-            x_src = x_node_dict[src_type]
-            x_dst = x_node_dict[dst_type]
-            alpha_src = (x_src * lin_src).sum(dim=-1)
-            alpha_dst = (x_dst * lin_dst).sum(dim=-1)
-            # propagate_type: (x_dst: PairTensor, alpha: PairTensor)
-            out = self.propagate(edge_index, x=(x_src, x_dst),
-                                 alpha=(alpha_src, alpha_dst), size=None)
-
-            out = F.relu(out)
-            out_dict[dst_type].append(out)
-
-        # iterate over node types:
-        semantic_attn_dict = {}
-        for node_type, outs in out_dict.items():
-            out, attn = group(outs, self.q, self.k_lin)
-            out_dict[node_type] = out
-            semantic_attn_dict[node_type] = attn
-
-        if return_semantic_attention_weights:
-            return out_dict, semantic_attn_dict
-
-        return out_dict
-
-    def message(self, x_j: Tensor, alpha_i: Tensor, alpha_j: Tensor,
-                index: Tensor, ptr: Optional[Tensor],
-                size_i: Optional[int]) -> Tensor:
-
-        alpha = alpha_j + alpha_i
-        alpha = F.leaky_relu(alpha, self.negative_slope)
-        alpha = softmax(alpha, index, ptr, size_i)
-        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
-        out = x_j * alpha.view(-1, self.heads, 1)
-        return out.view(-1, self.out_channels)
-
-    def __repr__(self) -> str:
-        return (f'{self.__class__.__name__}({self.out_channels}, '
-                f'heads={self.heads})')
+from typing import Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from torch_geometric.nn.conv import MessagePassing
+from torch_geometric.nn.dense import Linear
+from torch_geometric.nn.inits import glorot, reset
+from torch_geometric.typing import Adj, EdgeType, Metadata, NodeType, OptTensor
+from torch_geometric.utils import softmax
+
+
+def group(
+    xs: List[Tensor],
+    q: nn.Parameter,
+    k_lin: nn.Module,
+) -> Tuple[OptTensor, OptTensor]:
+
+    if len(xs) == 0:
+        return None, None
+    else:
+        num_edge_types = len(xs)
+        out = torch.stack(xs)
+        if out.numel() == 0:
+            return out.view(0, out.size(-1)), None
+        attn_score = (q * torch.tanh(k_lin(out)).mean(1)).sum(-1)
+        attn = F.softmax(attn_score, dim=0)
+        out = torch.sum(attn.view(num_edge_types, 1, -1) * out, dim=0)
+        return out, attn
+
+
+class HANConv(MessagePassing):
+    r"""
+    The Heterogenous Graph Attention Operator from the
+    `"Heterogenous Graph Attention Network"
+    <https://arxiv.org/pdf/1903.07293.pdf>`_ paper.
+
+    .. note::
+
+        For an example of using HANConv, see `examples/hetero/han_imdb.py
+        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
+        hetero/han_imdb.py>`_.
+
+    Args:
+        in_channels (int or Dict[str, int]): Size of each input sample of every
+            node type, or :obj:`-1` to derive the size from the first input(s)
+            to the forward method.
+        out_channels (int): Size of each output sample.
+        metadata (Tuple[List[str], List[Tuple[str, str, str]]]): The metadata
+            of the heterogeneous graph, *i.e.* its node and edge types given
+            by a list of strings and a list of string triplets, respectively.
+            See :meth:`torch_geometric.data.HeteroData.metadata` for more
+            information.
+        heads (int, optional): Number of multi-head-attentions.
+            (default: :obj:`1`)
+        negative_slope (float, optional): LeakyReLU angle of the negative
+            slope. (default: :obj:`0.2`)
+        dropout (float, optional): Dropout probability of the normalized
+            attention coefficients which exposes each node to a stochastically
+            sampled neighborhood during training. (default: :obj:`0`)
+        **kwargs (optional): Additional arguments of
+            :class:`torch_geometric.nn.conv.MessagePassing`.
+    """
+    def __init__(
+        self,
+        in_channels: Union[int, Dict[str, int]],
+        out_channels: int,
+        metadata: Metadata,
+        heads: int = 1,
+        negative_slope=0.2,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__(aggr='add', node_dim=0, **kwargs)
+
+        if not isinstance(in_channels, dict):
+            in_channels = {node_type: in_channels for node_type in metadata[0]}
+
+        self.heads = heads
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.negative_slope = negative_slope
+        self.metadata = metadata
+        self.dropout = dropout
+        self.k_lin = nn.Linear(out_channels, out_channels)
+        self.q = nn.Parameter(torch.empty(1, out_channels))
+
+        self.proj = nn.ModuleDict()
+        for node_type, in_channels in self.in_channels.items():
+            self.proj[node_type] = Linear(in_channels, out_channels)
+
+        self.lin_src = nn.ParameterDict()
+        self.lin_dst = nn.ParameterDict()
+        dim = out_channels // heads
+        for edge_type in metadata[1]:
+            edge_type = '__'.join(edge_type)
+            self.lin_src[edge_type] = nn.Parameter(torch.empty(1, heads, dim))
+            self.lin_dst[edge_type] = nn.Parameter(torch.empty(1, heads, dim))
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        super().reset_parameters()
+        reset(self.proj)
+        glorot(self.lin_src)
+        glorot(self.lin_dst)
+        self.k_lin.reset_parameters()
+        glorot(self.q)
+
+    def forward(
+        self,
+        x_dict: Dict[NodeType, Tensor],
+        edge_index_dict: Dict[EdgeType, Adj],
+        return_semantic_attention_weights: bool = False,
+    ) -> Union[Dict[NodeType, OptTensor], Tuple[Dict[NodeType, OptTensor],
+                                                Dict[NodeType, OptTensor]]]:
+        r"""Runs the forward pass of the module.
+
+        Args:
+            x_dict (Dict[str, torch.Tensor]): A dictionary holding node feature
+                information for each individual node type.
+            edge_index_dict (Dict[Tuple[str, str, str], torch.Tensor]): A
+                dictionary holding graph connectivity information for each
+                individual edge type, either as a :class:`torch.Tensor` of
+                shape :obj:`[2, num_edges]` or a
+                :class:`torch_sparse.SparseTensor`.
+            return_semantic_attention_weights (bool, optional): If set to
+                :obj:`True`, will additionally return the semantic-level
+                attention weights for each destination node type.
+                (default: :obj:`False`)
+        """
+        H, D = self.heads, self.out_channels // self.heads
+        x_node_dict, out_dict = {}, {}
+
+        # Iterate over node types:
+        for node_type, x in x_dict.items():
+            x_node_dict[node_type] = self.proj[node_type](x).view(-1, H, D)
+            out_dict[node_type] = []
+
+        # Iterate over edge types:
+        for edge_type, edge_index in edge_index_dict.items():
+            src_type, _, dst_type = edge_type
+            edge_type = '__'.join(edge_type)
+            lin_src = self.lin_src[edge_type]
+            lin_dst = self.lin_dst[edge_type]
+            x_src = x_node_dict[src_type]
+            x_dst = x_node_dict[dst_type]
+            alpha_src = (x_src * lin_src).sum(dim=-1)
+            alpha_dst = (x_dst * lin_dst).sum(dim=-1)
+            # propagate_type: (x_dst: PairTensor, alpha: PairTensor)
+            out = self.propagate(edge_index, x=(x_src, x_dst),
+                                 alpha=(alpha_src, alpha_dst), size=None)
+
+            out = F.relu(out)
+            out_dict[dst_type].append(out)
+
+        # iterate over node types:
+        semantic_attn_dict = {}
+        for node_type, outs in out_dict.items():
+            out, attn = group(outs, self.q, self.k_lin)
+            out_dict[node_type] = out
+            semantic_attn_dict[node_type] = attn
+
+        if return_semantic_attention_weights:
+            return out_dict, semantic_attn_dict
+
+        return out_dict
+
+    def message(self, x_j: Tensor, alpha_i: Tensor, alpha_j: Tensor,
+                index: Tensor, ptr: Optional[Tensor],
+                size_i: Optional[int]) -> Tensor:
+
+        alpha = alpha_j + alpha_i
+        alpha = F.leaky_relu(alpha, self.negative_slope)
+        alpha = softmax(alpha, index, ptr, size_i)
+        alpha = F.dropout(alpha, p=self.dropout, training=self.training)
+        out = x_j * alpha.view(-1, self.heads, 1)
+        return out.view(-1, self.out_channels)
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.out_channels}, '
+                f'heads={self.heads})')
diff --git a/torch_geometric/nn/models/attentive_fp.py b/torch_geometric/nn/models/attentive_fp.py
index 6fae299d0..c99b34c7b 100644
--- a/torch_geometric/nn/models/attentive_fp.py
+++ b/torch_geometric/nn/models/attentive_fp.py
@@ -51,8 +51,8 @@ class GATEConv(MessagePassing):
                 size_i: Optional[int]) -> Tensor:
 
         x_j = F.leaky_relu_(self.lin1(torch.cat([x_j, edge_attr], dim=-1)))
-        alpha_j = (x_j * self.att_l).sum(dim=-1)
-        alpha_i = (x_i * self.att_r).sum(dim=-1)
+        alpha_j = (x_j @ self.att_l.t()).squeeze(-1)
+        alpha_i = (x_i @ self.att_r.t()).squeeze(-1)
         alpha = alpha_j + alpha_i
         alpha = F.leaky_relu_(alpha)
         alpha = softmax(alpha, index, ptr, size_i)
diff --git a/torch_geometric/nn/models/rect.py b/torch_geometric/nn/models/rect.py
index 96d8f9cfe..d4e70c81d 100644
--- a/torch_geometric/nn/models/rect.py
+++ b/torch_geometric/nn/models/rect.py
@@ -1,153 +1,153 @@
-import copy
-
-import torch
-import torch.nn.functional as F
-from torch import Tensor
-from torch.nn import Linear
-
-from torch_geometric.nn import GCNConv
-from torch_geometric.typing import Adj, OptTensor, SparseTensor
-from torch_geometric.utils import scatter
-
-
-class RECT_L(torch.nn.Module):
-    r"""The RECT model, *i.e.* its supervised RECT-L part, from the
-    `"Network Embedding with Completely-imbalanced Labels"
-    <https://arxiv.org/abs/2007.03545>`_ paper.
-    In particular, a GCN model is trained that reconstructs semantic class
-    knowledge.
-
-    .. note::
-
-        For an example of using RECT, see `examples/rect.py
-        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
-        rect.py>`_.
-
-    Args:
-        in_channels (int): Size of each input sample.
-        hidden_channels (int): Intermediate size of each sample.
-        normalize (bool, optional): Whether to add self-loops and compute
-            symmetric normalization coefficients on-the-fly.
-            (default: :obj:`True`)
-        dropout (float, optional): The dropout probability.
-            (default: :obj:`0.0`)
-    """
-    def __init__(self, in_channels: int, hidden_channels: int,
-                 normalize: bool = True, dropout: float = 0.0):
-        super().__init__()
-        self.in_channels = in_channels
-        self.hidden_channels = hidden_channels
-        self.dropout = dropout
-
-        self.conv = GCNConv(in_channels, hidden_channels, normalize=normalize)
-        self.lin = Linear(hidden_channels, in_channels)
-
-        self.reset_parameters()
-
-    def reset_parameters(self):
-        r"""Resets all learnable parameters of the module."""
-        self.conv.reset_parameters()
-        self.lin.reset_parameters()
-        torch.nn.init.xavier_uniform_(self.lin.weight.data)
-
-    @torch.jit._overload_method
-    def forward(self, x, edge_index, edge_weight=None):
-        # type: (Tensor, SparseTensor, OptTensor) -> Tensor
-        pass
-
-    @torch.jit._overload_method
-    def forward(self, x, edge_index, edge_weight=None):
-        # type: (Tensor, Tensor, OptTensor) -> Tensor
-        pass
-
-    def forward(self, x: Tensor, edge_index: Adj,
-                edge_weight: OptTensor = None) -> Tensor:
-        """"""
-        x = self.conv(x, edge_index, edge_weight)
-        x = F.dropout(x, p=self.dropout, training=self.training)
-        return self.lin(x)
-
-    @torch.jit._overload_method
-    def embed(self, x, edge_index, edge_weight=None):
-        # type: (Tensor, SparseTensor, OptTensor) -> Tensor
-        pass
-
-    @torch.jit._overload_method
-    def embed(self, x, edge_index, edge_weight=None):
-        # type: (Tensor, Tensor, OptTensor) -> Tensor
-        pass
-
-    def embed(self, x: Tensor, edge_index: Adj,
-              edge_weight: OptTensor = None) -> Tensor:
-        with torch.no_grad():
-            return self.conv(x, edge_index, edge_weight)
-
-    def get_semantic_labels(self, x: Tensor, y: Tensor,
-                            mask: Tensor) -> Tensor:
-        r"""Replaces the original labels by their class-centers."""
-        with torch.no_grad():
-            y = y[mask]
-            mean = scatter(x[mask], y, dim=0, reduce='mean')
-            return mean[y]
-
-    def jittable(self, typing: str) -> torch.nn.Module:  # pragma: no cover
-        edge_index_type = typing.split(',')[1].strip()
-
-        class EdgeIndexJittable(torch.nn.Module):
-            def __init__(self, child: RECT_L):
-                super().__init__()
-                self.child = copy.deepcopy(child)
-                self.child.conv = self.child.conv.jittable()
-
-            def reset_parameters(self):
-                self.child.reset_parameters()
-
-            def forward(self, x: Tensor, edge_index: Tensor,
-                        edge_weight: OptTensor = None) -> Tensor:
-                return self.child(x, edge_index, edge_weight)
-
-            @torch.jit.export
-            def embed(self, x: Tensor, edge_index: Tensor,
-                      edge_weight: OptTensor = None) -> Tensor:
-                return self.child.embed(x, edge_index, edge_weight)
-
-            @torch.jit.export
-            def get_semantic_labels(self, x: Tensor, y: Tensor,
-                                    mask: Tensor) -> Tensor:
-                return self.child.get_semantic_labels(x, y, mask)
-
-        class SparseTensorJittable(torch.nn.Module):
-            def __init__(self, child: RECT_L):
-                super().__init__()
-                self.child = copy.deepcopy(child)
-                self.child.conv = self.child.conv.jittable()
-
-            def reset_parameters(self):
-                self.child.reset_parameters()
-
-            def forward(self, x: Tensor, edge_index: SparseTensor,
-                        edge_weight: OptTensor = None):
-                return self.child(x, edge_index, edge_weight)
-
-            @torch.jit.export
-            def embed(self, x: Tensor, edge_index: SparseTensor,
-                      edge_weight: OptTensor = None) -> Tensor:
-                return self.child.embed(x, edge_index, edge_weight)
-
-            @torch.jit.export
-            def get_semantic_labels(self, x: Tensor, y: Tensor,
-                                    mask: Tensor) -> Tensor:
-                return self.child.get_semantic_labels(x, y, mask)
-
-        if 'Tensor' == edge_index_type:
-            jittable_module = EdgeIndexJittable(self)
-        elif 'SparseTensor' == edge_index_type:
-            jittable_module = SparseTensorJittable(self)
-        else:
-            raise ValueError(f"Could not parse types '{typing}'")
-
-        return jittable_module
-
-    def __repr__(self) -> str:
-        return (f'{self.__class__.__name__}({self.in_channels}, '
-                f'{self.hidden_channels})')
+import copy
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.nn import Linear
+
+from torch_geometric.nn import GCNConv
+from torch_geometric.typing import Adj, OptTensor, SparseTensor
+from torch_geometric.utils import scatter
+
+
+class RECT_L(torch.nn.Module):
+    r"""The RECT model, *i.e.* its supervised RECT-L part, from the
+    `"Network Embedding with Completely-imbalanced Labels"
+    <https://arxiv.org/abs/2007.03545>`_ paper.
+    In particular, a GCN model is trained that reconstructs semantic class
+    knowledge.
+
+    .. note::
+
+        For an example of using RECT, see `examples/rect.py
+        <https://github.com/pyg-team/pytorch_geometric/blob/master/examples/
+        rect.py>`_.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Intermediate size of each sample.
+        normalize (bool, optional): Whether to add self-loops and compute
+            symmetric normalization coefficients on-the-fly.
+            (default: :obj:`True`)
+        dropout (float, optional): The dropout probability.
+            (default: :obj:`0.0`)
+    """
+    def __init__(self, in_channels: int, hidden_channels: int,
+                 normalize: bool = True, dropout: float = 0.0):
+        super().__init__()
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.dropout = dropout
+
+        self.conv = GCNConv(in_channels, hidden_channels, normalize=normalize)
+        self.lin = Linear(hidden_channels, in_channels)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        self.conv.reset_parameters()
+        self.lin.reset_parameters()
+        torch.nn.init.xavier_uniform_(self.lin.weight.data)
+
+    @torch.jit._overload_method
+    def forward(self, x, edge_index, edge_weight=None):
+        # type: (Tensor, SparseTensor, OptTensor) -> Tensor
+        pass
+
+    @torch.jit._overload_method
+    def forward(self, x, edge_index, edge_weight=None):
+        # type: (Tensor, Tensor, OptTensor) -> Tensor
+        pass
+
+    def forward(self, x: Tensor, edge_index: Adj,
+                edge_weight: OptTensor = None) -> Tensor:
+        """"""
+        x = self.conv(x, edge_index, edge_weight)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        return self.lin(x)
+
+    @torch.jit._overload_method
+    def embed(self, x, edge_index, edge_weight=None):
+        # type: (Tensor, SparseTensor, OptTensor) -> Tensor
+        pass
+
+    @torch.jit._overload_method
+    def embed(self, x, edge_index, edge_weight=None):
+        # type: (Tensor, Tensor, OptTensor) -> Tensor
+        pass
+
+    def embed(self, x: Tensor, edge_index: Adj,
+              edge_weight: OptTensor = None) -> Tensor:
+        with torch.no_grad():
+            return self.conv(x, edge_index, edge_weight)
+
+    def get_semantic_labels(self, x: Tensor, y: Tensor,
+                            mask: Tensor) -> Tensor:
+        r"""Replaces the original labels by their class-centers."""
+        with torch.no_grad():
+            y = y[mask]
+            mean = scatter(x[mask], y, dim=0, reduce='mean')
+            return mean[y]
+
+    def jittable(self, typing: str) -> torch.nn.Module:  # pragma: no cover
+        edge_index_type = typing.split(',')[1].strip()
+
+        class EdgeIndexJittable(torch.nn.Module):
+            def __init__(self, child: RECT_L):
+                super().__init__()
+                self.child = copy.deepcopy(child)
+                self.child.conv = self.child.conv.jittable()
+
+            def reset_parameters(self):
+                self.child.reset_parameters()
+
+            def forward(self, x: Tensor, edge_index: Tensor,
+                        edge_weight: OptTensor = None) -> Tensor:
+                return self.child(x, edge_index, edge_weight)
+
+            @torch.jit.export
+            def embed(self, x: Tensor, edge_index: Tensor,
+                      edge_weight: OptTensor = None) -> Tensor:
+                return self.child.embed(x, edge_index, edge_weight)
+
+            @torch.jit.export
+            def get_semantic_labels(self, x: Tensor, y: Tensor,
+                                    mask: Tensor) -> Tensor:
+                return self.child.get_semantic_labels(x, y, mask)
+
+        class SparseTensorJittable(torch.nn.Module):
+            def __init__(self, child: RECT_L):
+                super().__init__()
+                self.child = copy.deepcopy(child)
+                self.child.conv = self.child.conv.jittable()
+
+            def reset_parameters(self):
+                self.child.reset_parameters()
+
+            def forward(self, x: Tensor, edge_index: SparseTensor,
+                        edge_weight: OptTensor = None):
+                return self.child(x, edge_index, edge_weight)
+
+            @torch.jit.export
+            def embed(self, x: Tensor, edge_index: SparseTensor,
+                      edge_weight: OptTensor = None) -> Tensor:
+                return self.child.embed(x, edge_index, edge_weight)
+
+            @torch.jit.export
+            def get_semantic_labels(self, x: Tensor, y: Tensor,
+                                    mask: Tensor) -> Tensor:
+                return self.child.get_semantic_labels(x, y, mask)
+
+        if 'Tensor' == edge_index_type:
+            jittable_module = EdgeIndexJittable(self)
+        elif 'SparseTensor' == edge_index_type:
+            jittable_module = SparseTensorJittable(self)
+        else:
+            raise ValueError(f"Could not parse types '{typing}'")
+
+        return jittable_module
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.hidden_channels})')
diff --git a/torch_geometric/profile/__init__.py b/torch_geometric/profile/__init__.py
index 9f7340979..a72d953e6 100644
--- a/torch_geometric/profile/__init__.py
+++ b/torch_geometric/profile/__init__.py
@@ -4,6 +4,7 @@ from .profile import (
     print_time_total,
     rename_profile_file,
     torch_profile,
+    xpu_profile,
 )
 from .utils import count_parameters
 from .utils import get_model_size
@@ -21,6 +22,7 @@ __all__ = [
     'print_time_total',
     'rename_profile_file',
     'torch_profile',
+    'xpu_profile',
     'count_parameters',
     'get_model_size',
     'get_data_size',
diff --git a/torch_geometric/profile/profile.py b/torch_geometric/profile/profile.py
index fc8f0a887..748d7f57b 100644
--- a/torch_geometric/profile/profile.py
+++ b/torch_geometric/profile/profile.py
@@ -265,6 +265,15 @@ def torch_profile(export_chrome_trace=True, csv_data=None, write_csv=None):
         save_profile_data(csv_data, events, use_cuda)
 
 
+@contextmanager
+def xpu_profile(export_chrome_trace=True):
+    with torch.autograd.profiler_legacy.profile(use_xpu=True) as profile:
+        yield
+    print(profile.key_averages().table(sort_by='self_xpu_time_total'))
+    if export_chrome_trace:
+        profile.export_chrome_trace('timeline.json')
+
+
 def format_prof_time(time):
     # Profile time is in micro seconds, so format it appropriately:
     return round(time / 1e6, 3)
diff --git a/torch_geometric/transforms/remove_training_classes.py b/torch_geometric/transforms/remove_training_classes.py
index 1d29ff91c..7350bee75 100644
--- a/torch_geometric/transforms/remove_training_classes.py
+++ b/torch_geometric/transforms/remove_training_classes.py
@@ -1,27 +1,27 @@
-from typing import List
-
-from torch_geometric.data import Data
-from torch_geometric.data.datapipes import functional_transform
-from torch_geometric.transforms import BaseTransform
-
-
-@functional_transform('remove_training_classes')
-class RemoveTrainingClasses(BaseTransform):
-    r"""Removes classes from the node-level training set as given by
-    :obj:`data.train_mask`, *e.g.*, in order to get a zero-shot label scenario
-    (functional name: :obj:`remove_training_classes`).
-
-    Args:
-        classes (List[int]): The classes to remove from the training set.
-    """
-    def __init__(self, classes: List[int]):
-        self.classes = classes
-
-    def forward(self, data: Data) -> Data:
-        data.train_mask = data.train_mask.clone()
-        for i in self.classes:
-            data.train_mask[data.y == i] = False
-        return data
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}({self.classes})'
+from typing import List
+
+from torch_geometric.data import Data
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+
+
+@functional_transform('remove_training_classes')
+class RemoveTrainingClasses(BaseTransform):
+    r"""Removes classes from the node-level training set as given by
+    :obj:`data.train_mask`, *e.g.*, in order to get a zero-shot label scenario
+    (functional name: :obj:`remove_training_classes`).
+
+    Args:
+        classes (List[int]): The classes to remove from the training set.
+    """
+    def __init__(self, classes: List[int]):
+        self.classes = classes
+
+    def forward(self, data: Data) -> Data:
+        data.train_mask = data.train_mask.clone()
+        for i in self.classes:
+            data.train_mask[data.y == i] = False
+        return data
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.classes})'
diff --git a/torch_geometric/transforms/svd_feature_reduction.py b/torch_geometric/transforms/svd_feature_reduction.py
index 760d4e50d..c99798034 100644
--- a/torch_geometric/transforms/svd_feature_reduction.py
+++ b/torch_geometric/transforms/svd_feature_reduction.py
@@ -1,28 +1,28 @@
-import torch
-
-from torch_geometric.data import Data
-from torch_geometric.data.datapipes import functional_transform
-from torch_geometric.transforms import BaseTransform
-
-
-@functional_transform('svd_feature_reduction')
-class SVDFeatureReduction(BaseTransform):
-    r"""Dimensionality reduction of node features via Singular Value
-    Decomposition (SVD) (functional name: :obj:`svd_feature_reduction`).
-
-    Args:
-        out_channels (int): The dimensionlity of node features after
-            reduction.
-    """
-    def __init__(self, out_channels: int):
-        self.out_channels = out_channels
-
-    def forward(self, data: Data) -> Data:
-        if data.x.size(-1) > self.out_channels:
-            U, S, _ = torch.linalg.svd(data.x)
-            data.x = torch.mm(U[:, :self.out_channels],
-                              torch.diag(S[:self.out_channels]))
-        return data
-
-    def __repr__(self) -> str:
-        return f'{self.__class__.__name__}({self.out_channels})'
+import torch
+
+from torch_geometric.data import Data
+from torch_geometric.data.datapipes import functional_transform
+from torch_geometric.transforms import BaseTransform
+
+
+@functional_transform('svd_feature_reduction')
+class SVDFeatureReduction(BaseTransform):
+    r"""Dimensionality reduction of node features via Singular Value
+    Decomposition (SVD) (functional name: :obj:`svd_feature_reduction`).
+
+    Args:
+        out_channels (int): The dimensionlity of node features after
+            reduction.
+    """
+    def __init__(self, out_channels: int):
+        self.out_channels = out_channels
+
+    def forward(self, data: Data) -> Data:
+        if data.x.size(-1) > self.out_channels:
+            U, S, _ = torch.linalg.svd(data.x)
+            data.x = torch.mm(U[:, :self.out_channels],
+                              torch.diag(S[:self.out_channels]))
+        return data
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.out_channels})'
diff --git a/torch_geometric/utils/get_mesh_laplacian.py b/torch_geometric/utils/get_mesh_laplacian.py
index 98b8cba13..57a50884a 100644
--- a/torch_geometric/utils/get_mesh_laplacian.py
+++ b/torch_geometric/utils/get_mesh_laplacian.py
@@ -1,109 +1,109 @@
-from typing import Optional, Tuple
-
-import torch
-from torch import Tensor
-
-from torch_geometric.utils import add_self_loops, scatter, to_undirected
-
-
-def get_mesh_laplacian(
-    pos: Tensor,
-    face: Tensor,
-    normalization: Optional[str] = None,
-) -> Tuple[Tensor, Tensor]:
-    r"""Computes the mesh Laplacian of a mesh given by :obj:`pos` and
-    :obj:`face`. Computation is based on the cotangent matrix defined as
-
-    .. math::
-        \mathbf{C}_{ij} = \begin{cases}
-            \frac{\cot \angle_{ikj}~+\cot \angle_{ilj}}{2} &
-            \text{if } i, j \text{ is an edge} \\
-            -\sum_{j \in N(i)}{C_{ij}} &
-            \text{if } i \text{ is in the diagonal} \\
-            0 & \text{otherwise}
-      \end{cases}
-
-    Normalization depends on the mass matrix defined as
-
-    .. math::
-        \mathbf{M}_{ij} = \begin{cases}
-            a(i) & \text{if } i \text{ is in the diagonal} \\
-            0 & \text{otherwise}
-      \end{cases}
-
-    where :math:`a(i)` is obtained by joining the barycenters of the
-    triangles around vertex :math:`i`.
-
-    Args:
-        pos (Tensor): The node positions.
-        face (LongTensor): The face indices.
-        normalization (str, optional): The normalization scheme for the mesh
-            Laplacian (default: :obj:`None`):
-
-            1. :obj:`None`: No normalization
-            :math:`\mathbf{L} = \mathbf{C}`
-
-            2. :obj:`"sym"`: Symmetric normalization
-            :math:`\mathbf{L} = \mathbf{M}^{-1/2} \mathbf{C}\mathbf{M}^{-1/2}`
-
-            3. :obj:`"rw"`: Row-wise normalization
-            :math:`\mathbf{L} = \mathbf{M}^{-1} \mathbf{C}`
-    """
-    assert pos.size(1) == 3 and face.size(0) == 3
-
-    num_nodes = pos.shape[0]
-
-    def get_cots(left, centre, right):
-        left_pos, central_pos, right_pos = pos[left], pos[centre], pos[right]
-        left_vec = left_pos - central_pos
-        right_vec = right_pos - central_pos
-        dot = torch.einsum('ij, ij -> i', left_vec, right_vec)
-        cross = torch.norm(torch.cross(left_vec, right_vec, dim=1), dim=1)
-        cot = dot / cross  # cot = cos / sin
-        return cot / 2.0  # by definition
-
-    # For each triangle face, get all three cotangents:
-    cot_021 = get_cots(face[0], face[2], face[1])
-    cot_102 = get_cots(face[1], face[0], face[2])
-    cot_012 = get_cots(face[0], face[1], face[2])
-    cot_weight = torch.cat([cot_021, cot_102, cot_012])
-
-    # Face to edge:
-    cot_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
-    cot_index, cot_weight = to_undirected(cot_index, cot_weight)
-
-    # Compute the diagonal part:
-    cot_deg = scatter(cot_weight, cot_index[0], 0, num_nodes, reduce='sum')
-    edge_index, _ = add_self_loops(cot_index, num_nodes=num_nodes)
-    edge_weight = torch.cat([cot_weight, -cot_deg], dim=0)
-
-    if normalization is not None:
-
-        def get_areas(left, centre, right):
-            central_pos = pos[centre]
-            left_vec = pos[left] - central_pos
-            right_vec = pos[right] - central_pos
-            cross = torch.norm(torch.cross(left_vec, right_vec, dim=1), dim=1)
-            area = cross / 6.0  # one-third of a triangle's area is cross / 6.0
-            return area / 2.0  # since each corresponding area is counted twice
-
-        # Like before, but here we only need the diagonal (the mass matrix):
-        area_021 = get_areas(face[0], face[2], face[1])
-        area_102 = get_areas(face[1], face[0], face[2])
-        area_012 = get_areas(face[0], face[1], face[2])
-        area_weight = torch.cat([area_021, area_102, area_012])
-        area_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
-        area_index, area_weight = to_undirected(area_index, area_weight)
-        area_deg = scatter(area_weight, area_index[0], 0, num_nodes, 'sum')
-
-        if normalization == 'sym':
-            area_deg_inv_sqrt = area_deg.pow_(-0.5)
-            area_deg_inv_sqrt[area_deg_inv_sqrt == float('inf')] = 0.0
-            edge_weight = (area_deg_inv_sqrt[edge_index[0]] * edge_weight *
-                           area_deg_inv_sqrt[edge_index[1]])
-        elif normalization == 'rw':
-            area_deg_inv = 1.0 / area_deg
-            area_deg_inv[area_deg_inv == float('inf')] = 0.0
-            edge_weight = area_deg_inv[edge_index[0]] * edge_weight
-
-    return edge_index, edge_weight
+from typing import Optional, Tuple
+
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import add_self_loops, scatter, to_undirected
+
+
+def get_mesh_laplacian(
+    pos: Tensor,
+    face: Tensor,
+    normalization: Optional[str] = None,
+) -> Tuple[Tensor, Tensor]:
+    r"""Computes the mesh Laplacian of a mesh given by :obj:`pos` and
+    :obj:`face`. Computation is based on the cotangent matrix defined as
+
+    .. math::
+        \mathbf{C}_{ij} = \begin{cases}
+            \frac{\cot \angle_{ikj}~+\cot \angle_{ilj}}{2} &
+            \text{if } i, j \text{ is an edge} \\
+            -\sum_{j \in N(i)}{C_{ij}} &
+            \text{if } i \text{ is in the diagonal} \\
+            0 & \text{otherwise}
+      \end{cases}
+
+    Normalization depends on the mass matrix defined as
+
+    .. math::
+        \mathbf{M}_{ij} = \begin{cases}
+            a(i) & \text{if } i \text{ is in the diagonal} \\
+            0 & \text{otherwise}
+      \end{cases}
+
+    where :math:`a(i)` is obtained by joining the barycenters of the
+    triangles around vertex :math:`i`.
+
+    Args:
+        pos (Tensor): The node positions.
+        face (LongTensor): The face indices.
+        normalization (str, optional): The normalization scheme for the mesh
+            Laplacian (default: :obj:`None`):
+
+            1. :obj:`None`: No normalization
+            :math:`\mathbf{L} = \mathbf{C}`
+
+            2. :obj:`"sym"`: Symmetric normalization
+            :math:`\mathbf{L} = \mathbf{M}^{-1/2} \mathbf{C}\mathbf{M}^{-1/2}`
+
+            3. :obj:`"rw"`: Row-wise normalization
+            :math:`\mathbf{L} = \mathbf{M}^{-1} \mathbf{C}`
+    """
+    assert pos.size(1) == 3 and face.size(0) == 3
+
+    num_nodes = pos.shape[0]
+
+    def get_cots(left, centre, right):
+        left_pos, central_pos, right_pos = pos[left], pos[centre], pos[right]
+        left_vec = left_pos - central_pos
+        right_vec = right_pos - central_pos
+        dot = torch.einsum('ij, ij -> i', left_vec, right_vec)
+        cross = torch.norm(torch.cross(left_vec, right_vec, dim=1), dim=1)
+        cot = dot / cross  # cot = cos / sin
+        return cot / 2.0  # by definition
+
+    # For each triangle face, get all three cotangents:
+    cot_021 = get_cots(face[0], face[2], face[1])
+    cot_102 = get_cots(face[1], face[0], face[2])
+    cot_012 = get_cots(face[0], face[1], face[2])
+    cot_weight = torch.cat([cot_021, cot_102, cot_012])
+
+    # Face to edge:
+    cot_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
+    cot_index, cot_weight = to_undirected(cot_index, cot_weight)
+
+    # Compute the diagonal part:
+    cot_deg = scatter(cot_weight, cot_index[0], 0, num_nodes, reduce='sum')
+    edge_index, _ = add_self_loops(cot_index, num_nodes=num_nodes)
+    edge_weight = torch.cat([cot_weight, -cot_deg], dim=0)
+
+    if normalization is not None:
+
+        def get_areas(left, centre, right):
+            central_pos = pos[centre]
+            left_vec = pos[left] - central_pos
+            right_vec = pos[right] - central_pos
+            cross = torch.norm(torch.cross(left_vec, right_vec, dim=1), dim=1)
+            area = cross / 6.0  # one-third of a triangle's area is cross / 6.0
+            return area / 2.0  # since each corresponding area is counted twice
+
+        # Like before, but here we only need the diagonal (the mass matrix):
+        area_021 = get_areas(face[0], face[2], face[1])
+        area_102 = get_areas(face[1], face[0], face[2])
+        area_012 = get_areas(face[0], face[1], face[2])
+        area_weight = torch.cat([area_021, area_102, area_012])
+        area_index = torch.cat([face[:2], face[1:], face[::2]], dim=1)
+        area_index, area_weight = to_undirected(area_index, area_weight)
+        area_deg = scatter(area_weight, area_index[0], 0, num_nodes, 'sum')
+
+        if normalization == 'sym':
+            area_deg_inv_sqrt = area_deg.pow_(-0.5)
+            area_deg_inv_sqrt[area_deg_inv_sqrt == float('inf')] = 0.0
+            edge_weight = (area_deg_inv_sqrt[edge_index[0]] * edge_weight *
+                           area_deg_inv_sqrt[edge_index[1]])
+        elif normalization == 'rw':
+            area_deg_inv = 1.0 / area_deg
+            area_deg_inv[area_deg_inv == float('inf')] = 0.0
+            edge_weight = area_deg_inv[edge_index[0]] * edge_weight
+
+    return edge_index, edge_weight
diff --git a/torch_geometric/utils/spmm.py b/torch_geometric/utils/spmm.py
index 0e63e5312..b9dd77b99 100644
--- a/torch_geometric/utils/spmm.py
+++ b/torch_geometric/utils/spmm.py
@@ -46,7 +46,7 @@ def spmm(src: Adj, other: Tensor, reduce: str = "sum") -> Tensor:
         if (torch_geometric.typing.WITH_PT2 and other.dim() == 2
                 and not src.is_cuda() and not src.requires_grad()):
             # Use optimized PyTorch `torch.sparse.mm` path:
-            csr = src.to_torch_sparse_csr_tensor()
+            csr = src.to_torch_sparse_csr_tensor().to(other.dtype)
             return torch.sparse.mm(csr, other, reduce)
         return torch_sparse.matmul(src, other, reduce)
 


Test patch:
diff --git a/test/data/test_data.py b/test/data/test_data.py
index 4b4fd3ee1..2271fb59d 100644
--- a/test/data/test_data.py
+++ b/test/data/test_data.py
@@ -339,12 +339,13 @@ def test_data_share_memory():
 
     for data in data_list:
         assert not data.x.is_shared()
+        assert torch.all(data.x == 0.0)
 
     mp.spawn(run, args=(data_list, ), nprocs=4, join=True)
 
     for data in data_list:
         assert data.x.is_shared()
-        assert torch.allclose(data.x, torch.full((8, ), 4.))
+        assert torch.all(data.x > 0.0)
 
 
 def test_data_setter_properties():
diff --git a/test/nn/conv/test_han_conv.py b/test/nn/conv/test_han_conv.py
index 299d7e11f..caa7342d1 100644
--- a/test/nn/conv/test_han_conv.py
+++ b/test/nn/conv/test_han_conv.py
@@ -1,136 +1,136 @@
-import torch
-
-import torch_geometric.typing
-from torch_geometric.nn import HANConv
-from torch_geometric.typing import SparseTensor
-from torch_geometric.utils import coalesce, to_torch_csc_tensor
-
-
-def test_han_conv():
-    x_dict = {
-        'author': torch.randn(6, 16),
-        'paper': torch.randn(5, 12),
-        'term': torch.randn(4, 3)
-    }
-    edge_index1 = coalesce(torch.randint(0, 6, (2, 7)))
-    edge_index2 = coalesce(torch.randint(0, 5, (2, 4)))
-    edge_index3 = coalesce(torch.randint(0, 3, (2, 5)))
-    edge_index_dict = {
-        ('author', 'metapath0', 'author'): edge_index1,
-        ('paper', 'metapath1', 'paper'): edge_index2,
-        ('paper', 'metapath2', 'paper'): edge_index3,
-    }
-
-    adj_t_dict1 = {}
-    for edge_type, edge_index in edge_index_dict.items():
-        src_type, _, dst_type = edge_type
-        adj_t_dict1[edge_type] = to_torch_csc_tensor(
-            edge_index,
-            size=(x_dict[src_type].size(0), x_dict[dst_type].size(0)),
-        ).t()
-
-    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
-    in_channels = {'author': 16, 'paper': 12, 'term': 3}
-
-    conv = HANConv(in_channels, 16, metadata, heads=2)
-    assert str(conv) == 'HANConv(16, heads=2)'
-    out_dict1 = conv(x_dict, edge_index_dict)
-    assert len(out_dict1) == 3
-    assert out_dict1['author'].size() == (6, 16)
-    assert out_dict1['paper'].size() == (5, 16)
-    assert out_dict1['term'] is None
-    del out_dict1['term']
-    del x_dict['term']
-
-    out_dict2 = conv(x_dict, adj_t_dict1)
-    assert len(out_dict1) == len(out_dict2)
-    for key in out_dict1.keys():
-        assert torch.allclose(out_dict1[key], out_dict2[key], atol=1e-6)
-
-    if torch_geometric.typing.WITH_TORCH_SPARSE:
-        adj_t_dict2 = {}
-        for edge_type, edge_index in edge_index_dict.items():
-            adj_t_dict2[edge_type] = SparseTensor.from_edge_index(
-                edge_index,
-                sparse_sizes=adj_t_dict1[edge_type].size()[::-1],
-            ).t()
-        out_dict3 = conv(x_dict, adj_t_dict2)
-        assert len(out_dict1) == len(out_dict3)
-        for key in out_dict3.keys():
-            assert torch.allclose(out_dict1[key], out_dict3[key], atol=1e-6)
-
-    # Test non-zero dropout:
-    conv = HANConv(in_channels, 16, metadata, heads=2, dropout=0.1)
-    assert str(conv) == 'HANConv(16, heads=2)'
-    out_dict1 = conv(x_dict, edge_index_dict)
-    assert len(out_dict1) == 2
-    assert out_dict1['author'].size() == (6, 16)
-    assert out_dict1['paper'].size() == (5, 16)
-
-
-def test_han_conv_lazy():
-    x_dict = {
-        'author': torch.randn(6, 16),
-        'paper': torch.randn(5, 12),
-    }
-    edge_index1 = coalesce(torch.randint(0, 6, (2, 8)))
-    edge_index2 = coalesce(torch.randint(0, 5, (2, 6)))
-    edge_index_dict = {
-        ('author', 'to', 'author'): edge_index1,
-        ('paper', 'to', 'paper'): edge_index2,
-    }
-
-    adj_t_dict1 = {}
-    for edge_type, edge_index in edge_index_dict.items():
-        src_type, _, dst_type = edge_type
-        adj_t_dict1[edge_type] = to_torch_csc_tensor(
-            edge_index,
-            size=(x_dict[src_type].size(0), x_dict[dst_type].size(0)),
-        ).t()
-
-    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
-    conv = HANConv(-1, 16, metadata, heads=2)
-    assert str(conv) == 'HANConv(16, heads=2)'
-    out_dict1 = conv(x_dict, edge_index_dict)
-    assert len(out_dict1) == 2
-    assert out_dict1['author'].size() == (6, 16)
-    assert out_dict1['paper'].size() == (5, 16)
-
-    out_dict2 = conv(x_dict, adj_t_dict1)
-    assert len(out_dict1) == len(out_dict2)
-    for key in out_dict1.keys():
-        assert torch.allclose(out_dict1[key], out_dict2[key], atol=1e-6)
-
-    if torch_geometric.typing.WITH_TORCH_SPARSE:
-        adj_t_dict2 = {}
-        for edge_type, edge_index in edge_index_dict.items():
-            adj_t_dict2[edge_type] = SparseTensor.from_edge_index(
-                edge_index,
-                sparse_sizes=adj_t_dict1[edge_type].size()[::-1],
-            ).t()
-        out_dict3 = conv(x_dict, adj_t_dict2)
-        assert len(out_dict1) == len(out_dict3)
-        for key in out_dict1.keys():
-            assert torch.allclose(out_dict1[key], out_dict3[key], atol=1e-6)
-
-
-def test_han_conv_empty_tensor():
-    x_dict = {
-        'author': torch.randn(6, 16),
-        'paper': torch.empty(0, 12),
-    }
-    edge_index_dict = {
-        ('paper', 'to', 'author'): torch.empty((2, 0), dtype=torch.long),
-        ('author', 'to', 'paper'): torch.empty((2, 0), dtype=torch.long),
-        ('paper', 'to', 'paper'): torch.empty((2, 0), dtype=torch.long),
-    }
-
-    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
-    in_channels = {'author': 16, 'paper': 12}
-    conv = HANConv(in_channels, 16, metadata, heads=2)
-
-    out_dict = conv(x_dict, edge_index_dict)
-    assert len(out_dict) == 2
-    assert out_dict['author'].size() == (6, 16)
-    assert torch.all(out_dict['author'] == 0)
-    assert out_dict['paper'].size() == (0, 16)
+import torch
+
+import torch_geometric.typing
+from torch_geometric.nn import HANConv
+from torch_geometric.typing import SparseTensor
+from torch_geometric.utils import coalesce, to_torch_csc_tensor
+
+
+def test_han_conv():
+    x_dict = {
+        'author': torch.randn(6, 16),
+        'paper': torch.randn(5, 12),
+        'term': torch.randn(4, 3)
+    }
+    edge_index1 = coalesce(torch.randint(0, 6, (2, 7)))
+    edge_index2 = coalesce(torch.randint(0, 5, (2, 4)))
+    edge_index3 = coalesce(torch.randint(0, 3, (2, 5)))
+    edge_index_dict = {
+        ('author', 'metapath0', 'author'): edge_index1,
+        ('paper', 'metapath1', 'paper'): edge_index2,
+        ('paper', 'metapath2', 'paper'): edge_index3,
+    }
+
+    adj_t_dict1 = {}
+    for edge_type, edge_index in edge_index_dict.items():
+        src_type, _, dst_type = edge_type
+        adj_t_dict1[edge_type] = to_torch_csc_tensor(
+            edge_index,
+            size=(x_dict[src_type].size(0), x_dict[dst_type].size(0)),
+        ).t()
+
+    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
+    in_channels = {'author': 16, 'paper': 12, 'term': 3}
+
+    conv = HANConv(in_channels, 16, metadata, heads=2)
+    assert str(conv) == 'HANConv(16, heads=2)'
+    out_dict1 = conv(x_dict, edge_index_dict)
+    assert len(out_dict1) == 3
+    assert out_dict1['author'].size() == (6, 16)
+    assert out_dict1['paper'].size() == (5, 16)
+    assert out_dict1['term'] is None
+    del out_dict1['term']
+    del x_dict['term']
+
+    out_dict2 = conv(x_dict, adj_t_dict1)
+    assert len(out_dict1) == len(out_dict2)
+    for key in out_dict1.keys():
+        assert torch.allclose(out_dict1[key], out_dict2[key], atol=1e-6)
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj_t_dict2 = {}
+        for edge_type, edge_index in edge_index_dict.items():
+            adj_t_dict2[edge_type] = SparseTensor.from_edge_index(
+                edge_index,
+                sparse_sizes=adj_t_dict1[edge_type].size()[::-1],
+            ).t()
+        out_dict3 = conv(x_dict, adj_t_dict2)
+        assert len(out_dict1) == len(out_dict3)
+        for key in out_dict3.keys():
+            assert torch.allclose(out_dict1[key], out_dict3[key], atol=1e-6)
+
+    # Test non-zero dropout:
+    conv = HANConv(in_channels, 16, metadata, heads=2, dropout=0.1)
+    assert str(conv) == 'HANConv(16, heads=2)'
+    out_dict1 = conv(x_dict, edge_index_dict)
+    assert len(out_dict1) == 2
+    assert out_dict1['author'].size() == (6, 16)
+    assert out_dict1['paper'].size() == (5, 16)
+
+
+def test_han_conv_lazy():
+    x_dict = {
+        'author': torch.randn(6, 16),
+        'paper': torch.randn(5, 12),
+    }
+    edge_index1 = coalesce(torch.randint(0, 6, (2, 8)))
+    edge_index2 = coalesce(torch.randint(0, 5, (2, 6)))
+    edge_index_dict = {
+        ('author', 'to', 'author'): edge_index1,
+        ('paper', 'to', 'paper'): edge_index2,
+    }
+
+    adj_t_dict1 = {}
+    for edge_type, edge_index in edge_index_dict.items():
+        src_type, _, dst_type = edge_type
+        adj_t_dict1[edge_type] = to_torch_csc_tensor(
+            edge_index,
+            size=(x_dict[src_type].size(0), x_dict[dst_type].size(0)),
+        ).t()
+
+    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
+    conv = HANConv(-1, 16, metadata, heads=2)
+    assert str(conv) == 'HANConv(16, heads=2)'
+    out_dict1 = conv(x_dict, edge_index_dict)
+    assert len(out_dict1) == 2
+    assert out_dict1['author'].size() == (6, 16)
+    assert out_dict1['paper'].size() == (5, 16)
+
+    out_dict2 = conv(x_dict, adj_t_dict1)
+    assert len(out_dict1) == len(out_dict2)
+    for key in out_dict1.keys():
+        assert torch.allclose(out_dict1[key], out_dict2[key], atol=1e-6)
+
+    if torch_geometric.typing.WITH_TORCH_SPARSE:
+        adj_t_dict2 = {}
+        for edge_type, edge_index in edge_index_dict.items():
+            adj_t_dict2[edge_type] = SparseTensor.from_edge_index(
+                edge_index,
+                sparse_sizes=adj_t_dict1[edge_type].size()[::-1],
+            ).t()
+        out_dict3 = conv(x_dict, adj_t_dict2)
+        assert len(out_dict1) == len(out_dict3)
+        for key in out_dict1.keys():
+            assert torch.allclose(out_dict1[key], out_dict3[key], atol=1e-6)
+
+
+def test_han_conv_empty_tensor():
+    x_dict = {
+        'author': torch.randn(6, 16),
+        'paper': torch.empty(0, 12),
+    }
+    edge_index_dict = {
+        ('paper', 'to', 'author'): torch.empty((2, 0), dtype=torch.long),
+        ('author', 'to', 'paper'): torch.empty((2, 0), dtype=torch.long),
+        ('paper', 'to', 'paper'): torch.empty((2, 0), dtype=torch.long),
+    }
+
+    metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
+    in_channels = {'author': 16, 'paper': 12}
+    conv = HANConv(in_channels, 16, metadata, heads=2)
+
+    out_dict = conv(x_dict, edge_index_dict)
+    assert len(out_dict) == 2
+    assert out_dict['author'].size() == (6, 16)
+    assert torch.all(out_dict['author'] == 0)
+    assert out_dict['paper'].size() == (0, 16)
diff --git a/test/profile/test_profile.py b/test/profile/test_profile.py
index c4b5c38f3..b64d1125f 100644
--- a/test/profile/test_profile.py
+++ b/test/profile/test_profile.py
@@ -1,6 +1,7 @@
 import os.path
 import warnings
 
+import pytest
 import torch
 import torch.nn.functional as F
 
@@ -11,11 +12,12 @@ from torch_geometric.profile import (
     rename_profile_file,
     timeit,
 )
-from torch_geometric.profile.profile import torch_profile
+from torch_geometric.profile.profile import torch_profile, xpu_profile
 from torch_geometric.testing import (
     onlyCUDA,
     onlyLinux,
     onlyOnline,
+    onlyXPU,
     withCUDA,
     withPackage,
 )
@@ -105,3 +107,30 @@ def test_torch_profile(capfd, get_dataset, device):
     rename_profile_file('test_profile')
     assert os.path.exists('profile-test_profile.json')
     os.remove('profile-test_profile.json')
+
+
+@onlyXPU
+@onlyOnline
+@pytest.mark.parametrize('export_chrome_trace', [False, True])
+def test_xpu_profile(capfd, get_dataset, export_chrome_trace):
+    dataset = get_dataset(name='Cora')
+    device = torch.device('xpu')
+    data = dataset[0].to(device)
+    model = GraphSAGE(dataset.num_features, hidden_channels=64, num_layers=3,
+                      out_channels=dataset.num_classes).to(device)
+
+    with xpu_profile(export_chrome_trace):
+        model(data.x, data.edge_index)
+
+    out, _ = capfd.readouterr()
+    assert 'Self CPU' in out
+    if data.x.is_xpu:
+        assert 'Self XPU' in out
+
+    f_name = 'timeline.json'
+    f_exists = os.path.exists(f_name)
+    if not export_chrome_trace:
+        assert not f_exists
+    else:
+        assert f_exists
+        os.remove(f_name)
diff --git a/test/utils/test_get_mesh_laplacian.py b/test/utils/test_get_mesh_laplacian.py
index 04a6bdb83..de60ee04f 100644
--- a/test/utils/test_get_mesh_laplacian.py
+++ b/test/utils/test_get_mesh_laplacian.py
@@ -1,101 +1,101 @@
-import torch
-
-from torch_geometric.utils import get_mesh_laplacian
-
-
-def test_get_mesh_laplacian_of_cube():
-    pos = torch.tensor([
-        [1.0, 1.0, 1.0],
-        [1.0, -1.0, 1.0],
-        [-1.0, -1.0, 1.0],
-        [-1.0, 1.0, 1.0],
-        [1.0, 1.0, -1.0],
-        [1.0, -1.0, -1.0],
-        [-1.0, -1.0, -1.0],
-        [-1.0, 1.0, -1.0],
-    ])
-
-    face = torch.tensor([
-        [0, 1, 2],
-        [0, 3, 2],
-        [4, 5, 1],
-        [4, 0, 1],
-        [7, 6, 5],
-        [7, 4, 5],
-        [3, 2, 6],
-        [3, 7, 6],
-        [4, 0, 3],
-        [4, 7, 3],
-        [1, 5, 6],
-        [1, 2, 6],
-    ])
-
-    edge_index, edge_weight = get_mesh_laplacian(pos, face.t(),
-                                                 normalization='rw')
-
-    assert edge_index.tolist() == [
-        [
-            0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4,
-            4, 5, 5, 5, 5, 6, 6, 6, 6, 6, 7, 7, 7, 7, 0, 1, 2, 3, 4, 5, 6, 7
-        ],
-        [
-            1, 2, 3, 4, 0, 2, 4, 5, 6, 0, 1, 3, 6, 0, 2, 4, 6, 7, 0, 1, 3, 5,
-            7, 1, 4, 6, 7, 1, 2, 3, 5, 7, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7
-        ],
-    ]
-
-    assert torch.allclose(
-        edge_weight,
-        torch.tensor([
-            0.375, 0.0, 0.375, 0.375, 0.3, 0.3, 0.0, 0.3, 0.0, 0.0, 0.375,
-            0.375, 0.375, 0.3, 0.3, 0.0, 0.0, 0.3, 0.3, 0.0, 0.0, 0.3, 0.3,
-            0.375, 0.375, 0.375, 0.0, 0.0, 0.3, 0.0, 0.3, 0.3, 0.375, 0.375,
-            0.0, 0.375, -1.125, -0.9, -1.125, -0.9, -0.9, -1.125, -0.9, -1.125
-        ]))
-
-
-def test_get_mesh_laplacian_of_irregular_triangular_prism():
-    pos = torch.tensor([
-        [0.0, 0.0, 0.0],
-        [4.0, 0.0, 0.0],
-        [0.0, 0.0, -3.0],
-        [1.0, 5.0, -1.0],
-        [3.0, 5.0, -1.0],
-        [2.0, 5.0, -2.0],
-    ])
-
-    face = torch.tensor([
-        [0, 1, 2],
-        [3, 4, 5],
-        [0, 1, 4],
-        [0, 3, 4],
-        [1, 2, 5],
-        [1, 4, 5],
-        [2, 0, 3],
-        [2, 5, 3],
-    ])
-
-    edge_index, edge_weight = get_mesh_laplacian(pos, face.t(),
-                                                 normalization='rw')
-
-    assert edge_index.tolist() == [
-        [
-            0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5,
-            5, 5, 0, 1, 2, 3, 4, 5
-        ],
-        [
-            1, 2, 3, 4, 0, 2, 4, 5, 0, 1, 3, 5, 0, 2, 4, 5, 0, 1, 3, 5, 1, 2,
-            3, 4, 0, 1, 2, 3, 4, 5
-        ],
-    ]
-
-    assert torch.allclose(
-        edge_weight,
-        torch.tensor([
-            0.09730332, 0.15039921, 0.05081503, 0.00000000, 0.08726977,
-            0.03521059, 0.05363689, 0.00723919, 0.14497279, 0.03784235,
-            0.01629947, 0.03438699, 0.08362866, 0.02782887, 0.24252312,
-            0.40727590, 0.00000000, 0.08728313, 0.21507657, 0.38582093,
-            0.01117009, 0.04936920, 0.34247482, 0.36583540, -0.29851755,
-            -0.18335645, -0.23350160, -0.76125660, -0.68818060, -0.76884955
-        ]))
+import torch
+
+from torch_geometric.utils import get_mesh_laplacian
+
+
+def test_get_mesh_laplacian_of_cube():
+    pos = torch.tensor([
+        [1.0, 1.0, 1.0],
+        [1.0, -1.0, 1.0],
+        [-1.0, -1.0, 1.0],
+        [-1.0, 1.0, 1.0],
+        [1.0, 1.0, -1.0],
+        [1.0, -1.0, -1.0],
+        [-1.0, -1.0, -1.0],
+        [-1.0, 1.0, -1.0],
+    ])
+
+    face = torch.tensor([
+        [0, 1, 2],
+        [0, 3, 2],
+        [4, 5, 1],
+        [4, 0, 1],
+        [7, 6, 5],
+        [7, 4, 5],
+        [3, 2, 6],
+        [3, 7, 6],
+        [4, 0, 3],
+        [4, 7, 3],
+        [1, 5, 6],
+        [1, 2, 6],
+    ])
+
+    edge_index, edge_weight = get_mesh_laplacian(pos, face.t(),
+                                                 normalization='rw')
+
+    assert edge_index.tolist() == [
+        [
+            0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 3, 4, 4, 4, 4,
+            4, 5, 5, 5, 5, 6, 6, 6, 6, 6, 7, 7, 7, 7, 0, 1, 2, 3, 4, 5, 6, 7
+        ],
+        [
+            1, 2, 3, 4, 0, 2, 4, 5, 6, 0, 1, 3, 6, 0, 2, 4, 6, 7, 0, 1, 3, 5,
+            7, 1, 4, 6, 7, 1, 2, 3, 5, 7, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6, 7
+        ],
+    ]
+
+    assert torch.allclose(
+        edge_weight,
+        torch.tensor([
+            0.375, 0.0, 0.375, 0.375, 0.3, 0.3, 0.0, 0.3, 0.0, 0.0, 0.375,
+            0.375, 0.375, 0.3, 0.3, 0.0, 0.0, 0.3, 0.3, 0.0, 0.0, 0.3, 0.3,
+            0.375, 0.375, 0.375, 0.0, 0.0, 0.3, 0.0, 0.3, 0.3, 0.375, 0.375,
+            0.0, 0.375, -1.125, -0.9, -1.125, -0.9, -0.9, -1.125, -0.9, -1.125
+        ]))
+
+
+def test_get_mesh_laplacian_of_irregular_triangular_prism():
+    pos = torch.tensor([
+        [0.0, 0.0, 0.0],
+        [4.0, 0.0, 0.0],
+        [0.0, 0.0, -3.0],
+        [1.0, 5.0, -1.0],
+        [3.0, 5.0, -1.0],
+        [2.0, 5.0, -2.0],
+    ])
+
+    face = torch.tensor([
+        [0, 1, 2],
+        [3, 4, 5],
+        [0, 1, 4],
+        [0, 3, 4],
+        [1, 2, 5],
+        [1, 4, 5],
+        [2, 0, 3],
+        [2, 5, 3],
+    ])
+
+    edge_index, edge_weight = get_mesh_laplacian(pos, face.t(),
+                                                 normalization='rw')
+
+    assert edge_index.tolist() == [
+        [
+            0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5,
+            5, 5, 0, 1, 2, 3, 4, 5
+        ],
+        [
+            1, 2, 3, 4, 0, 2, 4, 5, 0, 1, 3, 5, 0, 2, 4, 5, 0, 1, 3, 5, 1, 2,
+            3, 4, 0, 1, 2, 3, 4, 5
+        ],
+    ]
+
+    assert torch.allclose(
+        edge_weight,
+        torch.tensor([
+            0.09730332, 0.15039921, 0.05081503, 0.00000000, 0.08726977,
+            0.03521059, 0.05363689, 0.00723919, 0.14497279, 0.03784235,
+            0.01629947, 0.03438699, 0.08362866, 0.02782887, 0.24252312,
+            0.40727590, 0.00000000, 0.08728313, 0.21507657, 0.38582093,
+            0.01117009, 0.04936920, 0.34247482, 0.36583540, -0.29851755,
+            -0.18335645, -0.23350160, -0.76125660, -0.68818060, -0.76884955
+        ]))
diff --git a/torch_geometric/testing/__init__.py b/torch_geometric/testing/__init__.py
index 83f982041..6a108d42f 100644
--- a/torch_geometric/testing/__init__.py
+++ b/torch_geometric/testing/__init__.py
@@ -4,6 +4,7 @@ from .decorators import (
     onlyLinux,
     onlyPython,
     onlyCUDA,
+    onlyXPU,
     onlyOnline,
     onlyGraphviz,
     onlyNeighborSampler,
@@ -22,6 +23,7 @@ __all__ = [
     'onlyLinux',
     'onlyPython',
     'onlyCUDA',
+    'onlyXPU',
     'onlyOnline',
     'onlyGraphviz',
     'onlyNeighborSampler',
diff --git a/torch_geometric/testing/decorators.py b/torch_geometric/testing/decorators.py
index b62625fa6..8c5e35c16 100644
--- a/torch_geometric/testing/decorators.py
+++ b/torch_geometric/testing/decorators.py
@@ -59,6 +59,20 @@ def onlyCUDA(func: Callable) -> Callable:
     )(func)
 
 
+def onlyXPU(func: Callable) -> Callable:
+    r"""A decorator to skip tests if XPU is not found."""
+    import pytest
+    try:
+        import intel_extension_for_pytorch as ipex
+        xpu_available = ipex.xpu.is_available()
+    except ImportError:
+        xpu_available = False
+    return pytest.mark.skipif(
+        not xpu_available,
+        reason="XPU not available",
+    )(func)
+
+
 def onlyOnline(func: Callable):
     r"""A decorator to skip tests if there exists no connection to the
     internet."""