pyg-nightly 2.7.0.dev20250701__py3-none-any.whl → 2.7.0.dev20250702__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250701.dist-info → pyg_nightly-2.7.0.dev20250702.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250701
+Version: 2.7.0.dev20250702
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>

{pyg_nightly-2.7.0.dev20250701.dist-info → pyg_nightly-2.7.0.dev20250702.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=DhH0xww6Nw5tVjgkem8O7_mt90VkyoLFJQzTSdB69n8,2250
+torch_geometric/__init__.py,sha256=ap-t4q8f9aTE0oAW_K5390u2Mlk8-S76rdeUEgPzglo,2250
 torch_geometric/_compile.py,sha256=9yqMTBKatZPr40WavJz9FjNi7pQj8YZAZOyZmmRGXgc,1351
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -295,7 +295,7 @@ torch_geometric/metrics/__init__.py,sha256=3krvDobW6vV5yHTjq2S2pmOXxNfysNG26muq7
 torch_geometric/metrics/link_pred.py,sha256=dtaI39JB-WqE1B-raiElns6xySRwmkbb9izbcyt6xHI,30886
 torch_geometric/nn/__init__.py,sha256=kQHHHUxFDht2ztD-XFQuv98TvC8MdodaFsIjAvltJBw,874
 torch_geometric/nn/data_parallel.py,sha256=YiybTWoSFyfSzlXAamZ_-y1f7B6tvDEFHOuy_AyJz9Q,4761
-torch_geometric/nn/encoding.py,sha256=QNjwWczYExZ1wRGBmpuqYbn6tB7NC4BU-DEgzjhcZqw,3115
+torch_geometric/nn/encoding.py,sha256=3DCOCO-XFt-lMb97sHWGN-4KeGUFY5lVo9P00SzrCNk,3559
 torch_geometric/nn/fx.py,sha256=PDtaHJAgodh4xf8FNl4fVxPGZJDbRaq3Q9z8qb1DNNI,16063
 torch_geometric/nn/glob.py,sha256=MdHjcUlHmFmTevzwND1_x7dXXJPzIDTBJRGOrGdZ8dQ,1088
 torch_geometric/nn/inits.py,sha256=_8FqacCLPz5Ft2zB5s6dtKGTKWtfrLyCLLuv1QvyKjk,2457
@@ -325,13 +325,13 @@ torch_geometric/nn/aggr/lcm.py,sha256=TcNqEvHnWpqOc9RFFioBAssQaUhOgMpH1_ovOmgl3w
 torch_geometric/nn/aggr/lstm.py,sha256=AdLa4rDd8t_X-GADDTOzRFuifSA0tIYVGKfoOckVtUE,2214
 torch_geometric/nn/aggr/mlp.py,sha256=sHQ4vQcZ-h2aOfFIBiXpAjr2lj7zHT3_TyqQr3WUjxQ,2514
 torch_geometric/nn/aggr/multi.py,sha256=theSIaDlLjGUyAtqDvOFORRpI9gYoZMXUtypX1PV5NQ,8170
-torch_geometric/nn/aggr/patch_transformer.py,sha256=SP--1IaXrHWjjGgH7yIPeO84b5NAwn65zHxaTid119o,5234
+torch_geometric/nn/aggr/patch_transformer.py,sha256=tWWBqBIuIPJfvFhkEs-S8cdEhuU1qxHsxoLh_ZnHznw,5498
 torch_geometric/nn/aggr/quantile.py,sha256=sRnKyt4CXr9RmjoPyTl4VUvXgSCMl9PG-fhCGsSZ76c,6189
 torch_geometric/nn/aggr/scaler.py,sha256=GV6gxUFBoKYMQTGybwzoPh708OY6k6chtUYmCIbFGXk,4638
 torch_geometric/nn/aggr/set2set.py,sha256=4GdmsjbBIrap3CG2naeFNsYe5eE-fhrNQOXM1-TIxyM,2446
 torch_geometric/nn/aggr/set_transformer.py,sha256=FG7_JizpFX14M6VSCwLSjYXYdJ1ZiQVbvnaYHIraiuM,4213
 torch_geometric/nn/aggr/sort.py,sha256=bvOOWnFkNOBOZih4rqVZQsjfeDX3vmXo1bpPSFD846w,2507
-torch_geometric/nn/aggr/utils.py,sha256=CLJ-ZrVWYIOBpdhQBLAz94dj3cMKKKc3qwGr4DFbiCU,8338
+torch_geometric/nn/aggr/utils.py,sha256=SQvdc0g6p_E2j0prA14MW2ekjEDvV-g545N0Q85uc-o,8625
 torch_geometric/nn/aggr/variance_preserving.py,sha256=fu-U_aGYpVLpgSFvVg0ONMe6nqoyv8tZ6Y35qMYTf9w,1126
 torch_geometric/nn/attention/__init__.py,sha256=wLKTmlfP7qL9sZHy4cmDFHEtdwa-MEKE1dT51L1_w10,192
 torch_geometric/nn/attention/performer.py,sha256=2PCDn4_-oNTao2-DkXIaoi18anP01OxRELF2pvp-jk8,7357
@@ -472,14 +472,14 @@ torch_geometric/nn/nlp/llm.py,sha256=DAv9jOZKXKQNVU2pNMyS1q8gVUtlin_unc6FjLhOYto
 torch_geometric/nn/nlp/sentence_transformer.py,sha256=q5M7SGtrUzoSiNhKCGFb7JatWiukdhNF6zdq2yiqxwE,4475
 torch_geometric/nn/nlp/vision_transformer.py,sha256=diVBefjIynzYs8WBlcpTeSVnw1PUecHY--B9Yd-W2hA,863
 torch_geometric/nn/norm/__init__.py,sha256=u2qIDrkbeuObGVXSAIftAlvSd6ouGTtxznCfD-59UiA,669
-torch_geometric/nn/norm/batch_norm.py,sha256=sJKrinHGwA-noIgteg1RD2W06rd0zskD-rXuY-36glY,8283
-torch_geometric/nn/norm/diff_group_norm.py,sha256=b57XvNekrUYGDjNJlGeqvaMGNJmHwopSF0_yyBWlLuA,4722
-torch_geometric/nn/norm/graph_norm.py,sha256=Tld_9_dzst4yEw58DZo4U--4QryA6pP2bsNfmqEDgrY,2727
+torch_geometric/nn/norm/batch_norm.py,sha256=fzUNmpdCUsMnNcso3PKDUdWc0UQvziK80-w0ZC0Vb8U,8706
+torch_geometric/nn/norm/diff_group_norm.py,sha256=mT0gM5a8txcAFNwZGKFu12qnNF5Pn95zrMx-RisRsh4,4938
+torch_geometric/nn/norm/graph_norm.py,sha256=VRmpi2jNYRQWXzX6Z0FmBxdEiV-EYXNPbGAGC0XNKH8,2964
 torch_geometric/nn/norm/graph_size_norm.py,sha256=sh5Nue1Ix2jC1T7o7KqOw0_TAOcpZ4VbYzhADWE97-M,1491
-torch_geometric/nn/norm/instance_norm.py,sha256=lUCZccuQNY8gfYUz-YRrNeSVckYuIJSFaW_m2HMp3iY,4685
-torch_geometric/nn/norm/layer_norm.py,sha256=m7a7Uoyx0zZhAxJ6Kj5N5DOg4zQnojO1FA1i661wW80,7835
+torch_geometric/nn/norm/instance_norm.py,sha256=L8VquSF7Jh5xfxA4YEcSO3IZ4fWR7VIiSukNeRXi4z0,4870
+torch_geometric/nn/norm/layer_norm.py,sha256=XiEyoXdDta6vlInLfwbJVsEgTkBFG6PJm6SpK99-cPE,8243
 torch_geometric/nn/norm/mean_subtraction_norm.py,sha256=KVHOp413mw7obwAN09Le6XdgobtCXpi4UKpjpG1M550,1322
-torch_geometric/nn/norm/msg_norm.py,sha256=zaQtqhs55LU-e6KPC4ylaSdge4KvEoseqOt7pmAzi2s,1662
+torch_geometric/nn/norm/msg_norm.py,sha256=NiV51ce1JgxVY6GbzktoSslDnZKWJrMJYZc_eBxz-pg,1903
 torch_geometric/nn/norm/pair_norm.py,sha256=IfHMiVYw_xsy035NakbPGdQVaVC-Ue3Oxwo651Vc47I,2824
 torch_geometric/nn/pool/__init__.py,sha256=VU9cPdLC-MPgt1kfS0ZwehfSD3g0V30VQuR1Wo0mzZE,14250
 torch_geometric/nn/pool/approx_knn.py,sha256=n7C8Cbar6o5tJcuAbzhM5hqMK26hW8dm5DopuocidO0,3967
@@ -640,7 +640,7 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=b-HnVesXjyJ_L1N-DnjiRiRVf7lhwKaBQF_2i5YMVSU,208
 torch_geometric/visualization/graph.py,sha256=mfZHXYfiU-CWMtfawYc80IxVwVmtK9hbIkSKhM_j7oI,14311
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250701.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250701.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250701.dist-info/METADATA,sha256=a6gNTxOjyu0eVS2qc87xRDkE8TutOw6BvtZGm15nUTI,63005
-pyg_nightly-2.7.0.dev20250701.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250702.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250702.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250702.dist-info/METADATA,sha256=66AyTfnfJvD0er8ePN_vOUgj6tD76JJy4QPaIvkh8bw,63005
+pyg_nightly-2.7.0.dev20250702.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250701'
+__version__ = '2.7.0.dev20250702'
 
 __all__ = [
     'Index',

torch_geometric/nn/aggr/patch_transformer.py

@@ -32,6 +32,8 @@ class PatchTransformerAggregation(Aggregation):
         aggr (str or list[str], optional): The aggregation module, *e.g.*,
             :obj:`"sum"`, :obj:`"mean"`, :obj:`"min"`, :obj:`"max"`,
             :obj:`"var"`, :obj:`"std"`. (default: :obj:`"mean"`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -43,6 +45,7 @@ class PatchTransformerAggregation(Aggregation):
         heads: int = 1,
         dropout: float = 0.0,
         aggr: Union[str, List[str]] = 'mean',
+        device: Optional[torch.device] = None,
     ) -> None:
         super().__init__()
 
@@ -55,12 +58,13 @@ class PatchTransformerAggregation(Aggregation):
         for aggr in self.aggrs:
             assert aggr in ['sum', 'mean', 'min', 'max', 'var', 'std']
 
-        self.lin = torch.nn.Linear(in_channels, hidden_channels)
+        self.lin = torch.nn.Linear(in_channels, hidden_channels, device=device)
         self.pad_projector = torch.nn.Linear(
             patch_size * hidden_channels,
             hidden_channels,
+            device=device,
         )
-        self.pe = PositionalEncoding(hidden_channels)
+        self.pe = PositionalEncoding(hidden_channels, device=device)
 
         self.blocks = torch.nn.ModuleList([
             MultiheadAttentionBlock(
@@ -68,12 +72,14 @@ class PatchTransformerAggregation(Aggregation):
                 heads=heads,
                 layer_norm=True,
                 dropout=dropout,
+                device=device,
             ) for _ in range(num_transformer_blocks)
         ])
 
         self.fc = torch.nn.Linear(
             hidden_channels * len(self.aggrs),
             out_channels,
+            device=device,
         )
 
     def reset_parameters(self) -> None:

torch_geometric/nn/aggr/utils.py

@@ -26,9 +26,11 @@ class MultiheadAttentionBlock(torch.nn.Module):
             normalization. (default: :obj:`True`)
         dropout (float, optional): Dropout probability of attention weights.
             (default: :obj:`0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(self, channels: int, heads: int = 1, layer_norm: bool = True,
-                 dropout: float = 0.0):
+                 dropout: float = 0.0, device: Optional[torch.device] = None):
         super().__init__()
 
         self.channels = channels
@@ -40,10 +42,13 @@ class MultiheadAttentionBlock(torch.nn.Module):
             heads,
             batch_first=True,
             dropout=dropout,
+            device=device,
         )
-        self.lin = Linear(channels, channels)
-        self.layer_norm1 = LayerNorm(channels) if layer_norm else None
-        self.layer_norm2 = LayerNorm(channels) if layer_norm else None
+        self.lin = Linear(channels, channels, device=device)
+        self.layer_norm1 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
+        self.layer_norm2 = LayerNorm(channels,
+                                     device=device) if layer_norm else None
 
     def reset_parameters(self):
         self.attn._reset_parameters()

torch_geometric/nn/encoding.py

@@ -1,4 +1,5 @@
 import math
+from typing import Optional
 
 import torch
 from torch import Tensor
@@ -23,12 +24,15 @@ class PositionalEncoding(torch.nn.Module):
         granularity (float, optional): The granularity of the positions. If
             set to smaller value, the encoder will capture more fine-grained
             changes in positions. (default: :obj:`1.0`)
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
         out_channels: int,
         base_freq: float = 1e-4,
         granularity: float = 1.0,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -40,7 +44,8 @@ class PositionalEncoding(torch.nn.Module):
         self.base_freq = base_freq
         self.granularity = granularity
 
-        frequency = torch.logspace(0, 1, out_channels // 2, base_freq)
+        frequency = torch.logspace(0, 1, out_channels // 2, base_freq,
+                                   device=device)
         self.register_buffer('frequency', frequency)
 
         self.reset_parameters()
@@ -75,13 +80,17 @@ class TemporalEncoding(torch.nn.Module):
 
     Args:
         out_channels (int): Size :math:`d` of each output sample.
+        device (torch.device, optional): The device of the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, out_channels: int):
+    def __init__(self, out_channels: int,
+                 device: Optional[torch.device] = None):
         super().__init__()
         self.out_channels = out_channels
 
         sqrt = math.sqrt(out_channels)
-        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels).view(1, -1)
+        weight = 1.0 / sqrt**torch.linspace(0, sqrt, out_channels,
+                                            device=device).view(1, -1)
         self.register_buffer('weight', weight)
 
         self.reset_parameters()

torch_geometric/nn/norm/batch_norm.py

@@ -39,6 +39,8 @@ class BatchNorm(torch.nn.Module):
             with only a single element will work as during in evaluation.
             That is the running mean and variance will be used.
             Requires :obj:`track_running_stats=True`. (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -48,6 +50,7 @@ class BatchNorm(torch.nn.Module):
         affine: bool = True,
         track_running_stats: bool = True,
         allow_single_element: bool = False,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -56,7 +59,7 @@ class BatchNorm(torch.nn.Module):
                              "'track_running_stats' to be set to `True`")
 
         self.module = torch.nn.BatchNorm1d(in_channels, eps, momentum, affine,
-                                           track_running_stats)
+                                           track_running_stats, device=device)
         self.in_channels = in_channels
         self.allow_single_element = allow_single_element
 
@@ -114,6 +117,8 @@ class HeteroBatchNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -123,6 +128,7 @@ class HeteroBatchNorm(torch.nn.Module):
         momentum: Optional[float] = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -134,17 +140,21 @@ class HeteroBatchNorm(torch.nn.Module):
         self.track_running_stats = track_running_stats
 
         if self.affine:
-            self.weight = Parameter(torch.empty(num_types, in_channels))
-            self.bias = Parameter(torch.empty(num_types, in_channels))
+            self.weight = Parameter(
+                torch.empty(num_types, in_channels, device=device))
+            self.bias = Parameter(
+                torch.empty(num_types, in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
 
         if self.track_running_stats:
-            self.register_buffer('running_mean',
-                                 torch.empty(num_types, in_channels))
-            self.register_buffer('running_var',
-                                 torch.empty(num_types, in_channels))
+            self.register_buffer(
+                'running_mean',
+                torch.empty(num_types, in_channels, device=device))
+            self.register_buffer(
+                'running_var',
+                torch.empty(num_types, in_channels, device=device))
             self.register_buffer('num_batches_tracked', torch.tensor(0))
         else:
             self.register_buffer('running_mean', None)

torch_geometric/nn/norm/diff_group_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 from torch import Tensor
 from torch.nn import BatchNorm1d, Linear
@@ -39,6 +41,8 @@ class DiffGroupNorm(torch.nn.Module):
             :obj:`False`, this module does not track such statistics and always
             uses batch statistics in both training and eval modes.
             (default: :obj:`True`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -49,6 +53,7 @@ class DiffGroupNorm(torch.nn.Module):
         momentum: float = 0.1,
         affine: bool = True,
         track_running_stats: bool = True,
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -56,9 +61,9 @@ class DiffGroupNorm(torch.nn.Module):
         self.groups = groups
         self.lamda = lamda
 
-        self.lin = Linear(in_channels, groups, bias=False)
+        self.lin = Linear(in_channels, groups, bias=False, device=device)
         self.norm = BatchNorm1d(groups * in_channels, eps, momentum, affine,
-                                track_running_stats)
+                                track_running_stats, device=device)
 
         self.reset_parameters()
 

torch_geometric/nn/norm/graph_norm.py

@@ -26,16 +26,21 @@ class GraphNorm(torch.nn.Module):
         in_channels (int): Size of each input sample.
         eps (float, optional): A value added to the denominator for numerical
             stability. (default: :obj:`1e-5`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, in_channels: int, eps: float = 1e-5):
+    def __init__(self, in_channels: int, eps: float = 1e-5,
+                 device: Optional[torch.device] = None):
         super().__init__()
 
         self.in_channels = in_channels
         self.eps = eps
 
-        self.weight = torch.nn.Parameter(torch.empty(in_channels))
-        self.bias = torch.nn.Parameter(torch.empty(in_channels))
-        self.mean_scale = torch.nn.Parameter(torch.empty(in_channels))
+        self.weight = torch.nn.Parameter(
+            torch.empty(in_channels, device=device))
+        self.bias = torch.nn.Parameter(torch.empty(in_channels, device=device))
+        self.mean_scale = torch.nn.Parameter(
+            torch.empty(in_channels, device=device))
 
         self.reset_parameters()
 

torch_geometric/nn/norm/instance_norm.py

@@ -1,5 +1,6 @@
 from typing import Optional
 
+import torch
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn.modules.instancenorm import _InstanceNorm
@@ -36,6 +37,8 @@ class InstanceNorm(_InstanceNorm):
             :obj:`False`, this module does not track such statistics and always
             uses instance statistics in both training and eval modes.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -44,9 +47,10 @@ class InstanceNorm(_InstanceNorm):
         momentum: float = 0.1,
         affine: bool = False,
         track_running_stats: bool = False,
+        device: Optional[torch.device] = None,
     ):
         super().__init__(in_channels, eps, momentum, affine,
-                         track_running_stats)
+                         track_running_stats, device=device)
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""

torch_geometric/nn/norm/layer_norm.py

@@ -35,6 +35,8 @@ class LayerNorm(torch.nn.Module):
             is used, each graph will be considered as an element to be
             normalized. If `"node"` is used, each node will be considered as
             an element to be normalized. (default: :obj:`"graph"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -42,6 +44,7 @@ class LayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'graph',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
 
@@ -51,8 +54,8 @@ class LayerNorm(torch.nn.Module):
         self.mode = mode
 
         if affine:
-            self.weight = Parameter(torch.empty(in_channels))
-            self.bias = Parameter(torch.empty(in_channels))
+            self.weight = Parameter(torch.empty(in_channels, device=device))
+            self.bias = Parameter(torch.empty(in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)
@@ -134,6 +137,8 @@ class HeteroLayerNorm(torch.nn.Module):
             normalization (:obj:`"node"`). If `"node"` is used, each node will
             be considered as an element to be normalized.
             (default: :obj:`"node"`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
     def __init__(
         self,
@@ -142,6 +147,7 @@ class HeteroLayerNorm(torch.nn.Module):
         eps: float = 1e-5,
         affine: bool = True,
         mode: str = 'node',
+        device: Optional[torch.device] = None,
     ):
         super().__init__()
         assert mode == 'node'
@@ -152,8 +158,10 @@ class HeteroLayerNorm(torch.nn.Module):
         self.affine = affine
 
         if affine:
-            self.weight = Parameter(torch.empty(num_types, in_channels))
-            self.bias = Parameter(torch.empty(num_types, in_channels))
+            self.weight = Parameter(
+                torch.empty(num_types, in_channels, device=device))
+            self.bias = Parameter(
+                torch.empty(num_types, in_channels, device=device))
         else:
             self.register_parameter('weight', None)
             self.register_parameter('bias', None)

torch_geometric/nn/norm/msg_norm.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import torch
 import torch.nn.functional as F
 from torch import Tensor
@@ -19,10 +21,14 @@ class MessageNorm(torch.nn.Module):
         learn_scale (bool, optional): If set to :obj:`True`, will learn the
             scaling factor :math:`s` of message normalization.
             (default: :obj:`False`)
+        device (torch.device, optional): The device to use for the module.
+            (default: :obj:`None`)
     """
-    def __init__(self, learn_scale: bool = False):
+    def __init__(self, learn_scale: bool = False,
+                 device: Optional[torch.device] = None):
         super().__init__()
-        self.scale = Parameter(torch.empty(1), requires_grad=learn_scale)
+        self.scale = Parameter(torch.empty(1, device=device),
+                               requires_grad=learn_scale)
         self.reset_parameters()
 
     def reset_parameters(self):