From 6b4c2db04cb000bb63c60046fb12c96674ffd8f2 Mon Sep 17 00:00:00 2001
From: rusty1s <matthias.fey@tu-dortmund.de>
Date: Wed, 15 Feb 2023 14:29:54 +0000
Subject: [PATCH] update

---
 examples/jit/linkx.py              | 52 ------------------------------
 torch_geometric/nn/models/linkx.py |  2 +-
 torch_geometric/nn/models/mlp.py   |  2 +-
 3 files changed, 2 insertions(+), 54 deletions(-)
 delete mode 100644 examples/jit/linkx.py

diff --git a/examples/jit/linkx.py b/examples/jit/linkx.py
deleted file mode 100644
index 05f7593a7458..000000000000
--- a/examples/jit/linkx.py
+++ /dev/null
@@ -1,52 +0,0 @@
-import os.path as osp
-
-import torch
-import torch.nn.functional as F
-
-from torch_geometric.datasets import LINKXDataset
-from torch_geometric.nn import LINKX
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-
-path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'LINKX')
-dataset = LINKXDataset(path, name='Penn94')
-# torch._C._jit_set_nvfuser_enabled(True)
-
-data = dataset[0].to(device)
-
-t = '(OptTensor, Tensor, OptTensor) -> Tensor'
-model = LINKX(data.num_nodes, data.num_features, hidden_channels=32,
-              out_channels=dataset.num_classes, num_layers=1,
-              num_edge_layers=1, num_node_layers=1,
-              dropout=0.5).jittable(t).to(device)
-model = torch.jit.script(model)
-optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=1e-3)
-
-
-def train():
-    model.train()
-    optimizer.zero_grad()
-    out = model(data.x, data.edge_index)
-    mask = data.train_mask[:, 0]  # Use the first set of the five masks.
-    loss = F.cross_entropy(out[mask], data.y[mask])
-    loss.backward()
-    optimizer.step()
-    return float(loss)
-
-
-@torch.no_grad()
-def test():
-    accs = []
-    model.eval()
-    pred = model(data.x, data.edge_index).argmax(dim=-1)
-    for _, mask in data('train_mask', 'val_mask', 'test_mask'):
-        mask = mask[:, 0]  # Use the first set of the five masks.
-        accs.append(int((pred[mask] == data.y[mask]).sum()) / int(mask.sum()))
-    return accs
-
-
-for epoch in range(1, 201):
-    loss = train()
-    train_acc, val_acc, test_acc = test()
-    print(f'Epoch: {epoch:03d}, Loss: {loss:.4f}, Train: {train_acc:.4f}, '
-          f'Val: {val_acc:.4f}, Test: {test_acc:.4f}')
diff --git a/torch_geometric/nn/models/linkx.py b/torch_geometric/nn/models/linkx.py
index f85dabc76e5f..5804f532c3ba 100644
--- 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
index 2f0db75a70b8..402a6c087363 100644
--- 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.