AssertionError when implementing heterogenous GNN

[PolarisRisingWar]

This is my output message:

Traceback (most recent call last):
  File "smalltraingraph.py", line 90, in <module>
    out = model(data.x_dict, data.edge_index_dict)
  File "myenvpath/lib/python3.8/site-packages/torch/fx/graph_module.py", line 630, in wrapped_call
    raise e.with_traceback(None)
AssertionError

I don't know why this happened? This is my data printed:

HeteroData(
  d1={ x=[96744, 768] },
  d2={ x=[94, 768] },
  d3={ x=[112, 768] },
  d4={ x=[11, 1] },
  (d1, similar1, d1)={ edge_index=[2, 252506] },
  (d2, similar2, d2)={ edge_index=[2, 134] },
  (d1, cites1, d2)={ edge_index=[2, 96744] },
  (d1, cites2, d3)={ edge_index=[2, 96744] },
  (d1, cites3, d4)={ edge_index=[2, 96744] }
)

But when using this code, I can run properly (I think I just have changed data object):

import torch
import torch_geometric.transforms as T
from torch_geometric.datasets import OGB_MAG
from torch_geometric.nn import SAGEConv, to_hetero


dataset = OGB_MAG(root='my_download/pyg_data', preprocess='metapath2vec', transform=T.ToUndirected())
data = dataset[0]

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x


model = GNN(hidden_channels=64, out_channels=8)
model = to_hetero(model, data.metadata(), aggr='sum')

with torch.no_grad():  # Initialize lazy modules.
    out = model(data.x_dict, data.edge_index_dict)

print(data)
print(out['paper'].size())

[rusty1s]

Sadly, the error messages produced by PyTorch ate not very descriptive here. My guess is that your data is incorrectly formatted. Try to run:

for edge_type in data.edge_types:
    src, _, dst = edge_type
    assert data[edge_type].edge_index[0].max() < data[src].num_nodes
    assert data[edge_type].edge_index[1].max() < data[dst].num_nodes

[PolarisRisingWar]

OK. After upgrading to nightly version, the bug messages become more clear. And I found that's because my edge index are in torch.int64 rather than torch.long, and my x are in torch.long and torch.float. I changed edge index to torch.long, and x to torch.float, which makes me able to run this code.

[rusty1s]

Super :)