update

pyg-team · Feb 15, 2023 · 6b4c2db · 6b4c2db
1 parent ad121a9
commit 6b4c2db
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 54 deletions.
diff --git a/examples/jit/linkx.py b/examples/jit/linkx.py
diff --git a/torch_geometric/nn/models/linkx.py b/torch_geometric/nn/models/linkx.py
@@ -120,7 +120,7 @@ def __init__(
 
         if self.num_edge_layers > 1:
             self.edge_norm = BatchNorm1d(hidden_channels)
-            channels = [hidden_channels] * self.num_edge_layers
+            channels = [hidden_channels] * num_edge_layers
             self.edge_mlp = MLP(channels, dropout=0., act_first=True)
         else:
             self.edge_norm = None

diff --git a/torch_geometric/nn/models/mlp.py b/torch_geometric/nn/models/mlp.py
@@ -189,7 +189,7 @@ def forward(
         self,
         x: Tensor,
         return_emb: NoneType = None,
-    ):
+    ) -> Tensor:
         r"""
         Args:
             x (torch.Tensor): The source tensor.