prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Computing the size of the in-degree histogram tensor dynamically in the PNA example
  ### 🛠 Proposed Refactor
  
  The `deg` tensor in the PNA example is initialized with the size `5`. 
  https://github.com/pyg-team/pytorch_geometric/blob/50b7bfc4a59b5b6f7ec547ff862985f3b2e22798/examples/pna.py#L23
  
  This value will obviously be different for different datasets. One can iterate over the training data and compute the maximum degree any node has. Then, the histogram tensor can be initialized with that value. Something like this:
  ```python
  # compute the maximum in-degree in the training data
  max_degree = 0
  for data in train_dataset:
      d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
      if d.max().item() > max_degree:
          max_degree = d.max().item()
  
  # create the in-degree histogram tensor
  deg = torch.zeros(max_degree + 1, dtype=torch.long)
  for data in train_dataset:
      d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
      deg += torch.bincount(d, minlength=deg.numel())
  ``` 
  
  ### Suggest a potential alternative/fix
  
  The results of the `degree` function can also be cached to avoid iterating two times. Further, for custom datasets where the split is made with indices or boolean masks e.g `Data(x=[10000, 100], edge_index=[2, 200000], edge_attr=[200000, 20], y=[10000], train_mask=[10000], val_mask=[10000], test_mask=[10000])`, the `subgraph` utility can be used.
  ```python
  tr_subgraph = data.subgraph(data.train_mask)
  # compute the in-degree of all the training nodes
  d = degree(index=tr_subgraph.edge_index[1], num_nodes=tr_subgraph.num_nodes, dtype=torch.long)
  
  # get the maximum in-degree, this will be the size of the histogram tensor
  max_degree = d.max().item()
  
  # create the in-degree histogram tensor
  deg = torch.zeros(max_degree + 1, dtype=torch.long)
  deg += torch.bincount(d, minlength=deg.numel())
  ```
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `examples/pna.py`
  Content:
  ```
  1 import os.path as osp
  2 
  3 import torch
  4 import torch.nn.functional as F
  5 from torch.nn import Embedding, Linear, ModuleList, ReLU, Sequential
  6 from torch.optim.lr_scheduler import ReduceLROnPlateau
  7 
  8 from torch_geometric.datasets import ZINC
  9 from torch_geometric.loader import DataLoader
  10 from torch_geometric.nn import BatchNorm, PNAConv, global_add_pool
  11 from torch_geometric.utils import degree
  12 
  13 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ZINC')
  14 train_dataset = ZINC(path, subset=True, split='train')
  15 val_dataset = ZINC(path, subset=True, split='val')
  16 test_dataset = ZINC(path, subset=True, split='test')
  17 
  18 train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)
  19 val_loader = DataLoader(val_dataset, batch_size=128)
  20 test_loader = DataLoader(test_dataset, batch_size=128)
  21 
  22 # Compute in-degree histogram over training data.
  23 deg = torch.zeros(5, dtype=torch.long)
  24 for data in train_dataset:
  25     d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
  26     deg += torch.bincount(d, minlength=deg.numel())
  27 
  28 
  29 class Net(torch.nn.Module):
  30     def __init__(self):
  31         super().__init__()
  32 
  33         self.node_emb = Embedding(21, 75)
  34         self.edge_emb = Embedding(4, 50)
  35 
  36         aggregators = ['mean', 'min', 'max', 'std']
  37         scalers = ['identity', 'amplification', 'attenuation']
  38 
  39         self.convs = ModuleList()
  40         self.batch_norms = ModuleList()
  41         for _ in range(4):
  42             conv = PNAConv(in_channels=75, out_channels=75,
  43                            aggregators=aggregators, scalers=scalers, deg=deg,
  44                            edge_dim=50, towers=5, pre_layers=1, post_layers=1,
  45                            divide_input=False)
  46             self.convs.append(conv)
  47             self.batch_norms.append(BatchNorm(75))
  48 
  49         self.mlp = Sequential(Linear(75, 50), ReLU(), Linear(50, 25), ReLU(),
  50                               Linear(25, 1))
  51 
  52     def forward(self, x, edge_index, edge_attr, batch):
  53         x = self.node_emb(x.squeeze())
  54         edge_attr = self.edge_emb(edge_attr)
  55 
  56         for conv, batch_norm in zip(self.convs, self.batch_norms):
  57             x = F.relu(batch_norm(conv(x, edge_index, edge_attr)))
  58 
  59         x = global_add_pool(x, batch)
  60         return self.mlp(x)
  61 
  62 
  63 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
  64 model = Net().to(device)
  65 optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
  66 scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=20,
  67                               min_lr=0.00001)
  68 
  69 
  70 def train(epoch):
  71     model.train()
  72 
  73     total_loss = 0
  74     for data in train_loader:
  75         data = data.to(device)
  76         optimizer.zero_grad()
  77         out = model(data.x, data.edge_index, data.edge_attr, data.batch)
  78         loss = (out.squeeze() - data.y).abs().mean()
  79         loss.backward()
  80         total_loss += loss.item() * data.num_graphs
  81         optimizer.step()
  82     return total_loss / len(train_loader.dataset)
  83 
  84 
  85 @torch.no_grad()
  86 def test(loader):
  87     model.eval()
  88 
  89     total_error = 0
  90     for data in loader:
  91         data = data.to(device)
  92         out = model(data.x, data.edge_index, data.edge_attr, data.batch)
  93         total_error += (out.squeeze() - data.y).abs().sum().item()
  94     return total_error / len(loader.dataset)
  95 
  96 
  97 for epoch in range(1, 301):
  98     loss = train(epoch)
  99     val_mae = test(val_loader)
  100     test_mae = test(test_loader)
  101     scheduler.step(val_mae)
  102     print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Val: {val_mae:.4f}, '
  103           f'Test: {test_mae:.4f}')
  104
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/examples/pna.py b/examples/pna.py
  --- a/examples/pna.py
  +++ b/examples/pna.py
  @@ -19,8 +19,14 @@
   val_loader = DataLoader(val_dataset, batch_size=128)
   test_loader = DataLoader(test_dataset, batch_size=128)
   
  -# Compute in-degree histogram over training data.
  -deg = torch.zeros(5, dtype=torch.long)
  +# Compute the maximum in-degree in the training data.
  +max_degree = -1
  +for data in train_dataset:
  +    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
  +    max_degree = max(max_degree, int(d.max()))
  +
  +# Compute the in-degree histogram tensor
  +deg = torch.zeros(max_degree + 1, dtype=torch.long)
   for data in train_dataset:
       d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)
       deg += torch.bincount(d, minlength=deg.numel())
  
verification_info: {"golden_diff": "diff --git a/examples/pna.py b/examples/pna.py\n--- a/examples/pna.py\n+++ b/examples/pna.py\n@@ -19,8 +19,14 @@\n val_loader = DataLoader(val_dataset, batch_size=128)\n test_loader = DataLoader(test_dataset, batch_size=128)\n \n-# Compute in-degree histogram over training data.\n-deg = torch.zeros(5, dtype=torch.long)\n+# Compute the maximum in-degree in the training data.\n+max_degree = -1\n+for data in train_dataset:\n+    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\n+    max_degree = max(max_degree, int(d.max()))\n+\n+# Compute the in-degree histogram tensor\n+deg = torch.zeros(max_degree + 1, dtype=torch.long)\n for data in train_dataset:\n     d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\n     deg += torch.bincount(d, minlength=deg.numel())\n", "issue": "Computing the size of the in-degree histogram tensor dynamically in the PNA example\n### \ud83d\udee0 Proposed Refactor\n\nThe `deg` tensor in the PNA example is initialized with the size `5`. \r\nhttps://github.com/pyg-team/pytorch_geometric/blob/50b7bfc4a59b5b6f7ec547ff862985f3b2e22798/examples/pna.py#L23\r\n\r\nThis value will obviously be different for different datasets. One can iterate over the training data and compute the maximum degree any node has. Then, the histogram tensor can be initialized with that value. Something like this:\r\n```python\r\n# compute the maximum in-degree in the training data\r\nmax_degree = 0\r\nfor data in train_dataset:\r\n    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\r\n    if d.max().item() > max_degree:\r\n        max_degree = d.max().item()\r\n\r\n# create the in-degree histogram tensor\r\ndeg = torch.zeros(max_degree + 1, dtype=torch.long)\r\nfor data in train_dataset:\r\n    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\r\n    deg += torch.bincount(d, minlength=deg.numel())\r\n``` \n\n### Suggest a potential alternative/fix\n\nThe results of the `degree` function can also be cached to avoid iterating two times. Further, for custom datasets where the split is made with indices or boolean masks e.g `Data(x=[10000, 100], edge_index=[2, 200000], edge_attr=[200000, 20], y=[10000], train_mask=[10000], val_mask=[10000], test_mask=[10000])`, the `subgraph` utility can be used.\r\n```python\r\ntr_subgraph = data.subgraph(data.train_mask)\r\n# compute the in-degree of all the training nodes\r\nd = degree(index=tr_subgraph.edge_index[1], num_nodes=tr_subgraph.num_nodes, dtype=torch.long)\r\n\r\n# get the maximum in-degree, this will be the size of the histogram tensor\r\nmax_degree = d.max().item()\r\n\r\n# create the in-degree histogram tensor\r\ndeg = torch.zeros(max_degree + 1, dtype=torch.long)\r\ndeg += torch.bincount(d, minlength=deg.numel())\r\n```\n", "before_files": [{"content": "import os.path as osp\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.nn import Embedding, Linear, ModuleList, ReLU, Sequential\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nfrom torch_geometric.datasets import ZINC\nfrom torch_geometric.loader import DataLoader\nfrom torch_geometric.nn import BatchNorm, PNAConv, global_add_pool\nfrom torch_geometric.utils import degree\n\npath = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ZINC')\ntrain_dataset = ZINC(path, subset=True, split='train')\nval_dataset = ZINC(path, subset=True, split='val')\ntest_dataset = ZINC(path, subset=True, split='test')\n\ntrain_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=128)\ntest_loader = DataLoader(test_dataset, batch_size=128)\n\n# Compute in-degree histogram over training data.\ndeg = torch.zeros(5, dtype=torch.long)\nfor data in train_dataset:\n    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\n    deg += torch.bincount(d, minlength=deg.numel())\n\n\nclass Net(torch.nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.node_emb = Embedding(21, 75)\n        self.edge_emb = Embedding(4, 50)\n\n        aggregators = ['mean', 'min', 'max', 'std']\n        scalers = ['identity', 'amplification', 'attenuation']\n\n        self.convs = ModuleList()\n        self.batch_norms = ModuleList()\n        for _ in range(4):\n            conv = PNAConv(in_channels=75, out_channels=75,\n                           aggregators=aggregators, scalers=scalers, deg=deg,\n                           edge_dim=50, towers=5, pre_layers=1, post_layers=1,\n                           divide_input=False)\n            self.convs.append(conv)\n            self.batch_norms.append(BatchNorm(75))\n\n        self.mlp = Sequential(Linear(75, 50), ReLU(), Linear(50, 25), ReLU(),\n                              Linear(25, 1))\n\n    def forward(self, x, edge_index, edge_attr, batch):\n        x = self.node_emb(x.squeeze())\n        edge_attr = self.edge_emb(edge_attr)\n\n        for conv, batch_norm in zip(self.convs, self.batch_norms):\n            x = F.relu(batch_norm(conv(x, edge_index, edge_attr)))\n\n        x = global_add_pool(x, batch)\n        return self.mlp(x)\n\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodel = Net().to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)\nscheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=20,\n                              min_lr=0.00001)\n\n\ndef train(epoch):\n    model.train()\n\n    total_loss = 0\n    for data in train_loader:\n        data = data.to(device)\n        optimizer.zero_grad()\n        out = model(data.x, data.edge_index, data.edge_attr, data.batch)\n        loss = (out.squeeze() - data.y).abs().mean()\n        loss.backward()\n        total_loss += loss.item() * data.num_graphs\n        optimizer.step()\n    return total_loss / len(train_loader.dataset)\n\n\n@torch.no_grad()\ndef test(loader):\n    model.eval()\n\n    total_error = 0\n    for data in loader:\n        data = data.to(device)\n        out = model(data.x, data.edge_index, data.edge_attr, data.batch)\n        total_error += (out.squeeze() - data.y).abs().sum().item()\n    return total_error / len(loader.dataset)\n\n\nfor epoch in range(1, 301):\n    loss = train(epoch)\n    val_mae = test(val_loader)\n    test_mae = test(test_loader)\n    scheduler.step(val_mae)\n    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Val: {val_mae:.4f}, '\n          f'Test: {test_mae:.4f}')\n", "path": "examples/pna.py"}], "after_files": [{"content": "import os.path as osp\n\nimport torch\nimport torch.nn.functional as F\nfrom torch.nn import Embedding, Linear, ModuleList, ReLU, Sequential\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\n\nfrom torch_geometric.datasets import ZINC\nfrom torch_geometric.loader import DataLoader\nfrom torch_geometric.nn import BatchNorm, PNAConv, global_add_pool\nfrom torch_geometric.utils import degree\n\npath = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'ZINC')\ntrain_dataset = ZINC(path, subset=True, split='train')\nval_dataset = ZINC(path, subset=True, split='val')\ntest_dataset = ZINC(path, subset=True, split='test')\n\ntrain_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=128)\ntest_loader = DataLoader(test_dataset, batch_size=128)\n\n# Compute the maximum in-degree in the training data.\nmax_degree = -1\nfor data in train_dataset:\n    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\n    max_degree = max(max_degree, int(d.max()))\n\n# Compute the in-degree histogram tensor\ndeg = torch.zeros(max_degree + 1, dtype=torch.long)\nfor data in train_dataset:\n    d = degree(data.edge_index[1], num_nodes=data.num_nodes, dtype=torch.long)\n    deg += torch.bincount(d, minlength=deg.numel())\n\n\nclass Net(torch.nn.Module):\n    def __init__(self):\n        super().__init__()\n\n        self.node_emb = Embedding(21, 75)\n        self.edge_emb = Embedding(4, 50)\n\n        aggregators = ['mean', 'min', 'max', 'std']\n        scalers = ['identity', 'amplification', 'attenuation']\n\n        self.convs = ModuleList()\n        self.batch_norms = ModuleList()\n        for _ in range(4):\n            conv = PNAConv(in_channels=75, out_channels=75,\n                           aggregators=aggregators, scalers=scalers, deg=deg,\n                           edge_dim=50, towers=5, pre_layers=1, post_layers=1,\n                           divide_input=False)\n            self.convs.append(conv)\n            self.batch_norms.append(BatchNorm(75))\n\n        self.mlp = Sequential(Linear(75, 50), ReLU(), Linear(50, 25), ReLU(),\n                              Linear(25, 1))\n\n    def forward(self, x, edge_index, edge_attr, batch):\n        x = self.node_emb(x.squeeze())\n        edge_attr = self.edge_emb(edge_attr)\n\n        for conv, batch_norm in zip(self.convs, self.batch_norms):\n            x = F.relu(batch_norm(conv(x, edge_index, edge_attr)))\n\n        x = global_add_pool(x, batch)\n        return self.mlp(x)\n\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nmodel = Net().to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)\nscheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=20,\n                              min_lr=0.00001)\n\n\ndef train(epoch):\n    model.train()\n\n    total_loss = 0\n    for data in train_loader:\n        data = data.to(device)\n        optimizer.zero_grad()\n        out = model(data.x, data.edge_index, data.edge_attr, data.batch)\n        loss = (out.squeeze() - data.y).abs().mean()\n        loss.backward()\n        total_loss += loss.item() * data.num_graphs\n        optimizer.step()\n    return total_loss / len(train_loader.dataset)\n\n\n@torch.no_grad()\ndef test(loader):\n    model.eval()\n\n    total_error = 0\n    for data in loader:\n        data = data.to(device)\n        out = model(data.x, data.edge_index, data.edge_attr, data.batch)\n        total_error += (out.squeeze() - data.y).abs().sum().item()\n    return total_error / len(loader.dataset)\n\n\nfor epoch in range(1, 301):\n    loss = train(epoch)\n    val_mae = test(val_loader)\n    test_mae = test(test_loader)\n    scheduler.step(val_mae)\n    print(f'Epoch: {epoch:02d}, Loss: {loss:.4f}, Val: {val_mae:.4f}, '\n          f'Test: {test_mae:.4f}')\n", "path": "examples/pna.py"}]}