Problem with TransformerConv in uninitialized submodule and hetero graphs

[chrisn-pik]

🐛 Describe the bug

When using heterogeneous graphs and using TransformerConv in a submodul I get an error that indices in 'edge_index' are larger and invalid.

To reproduce this issue:

from torch_geometric.nn import TransformerConv
from torch_geometric.nn import to_hetero
import torch
import torch.nn as nn
from torch_geometric.datasets import FakeHeteroDataset

dataset = FakeHeteroDataset(1, edge_dim=3)
data = dataset[0]

class TransformerObject(nn.Module):
    def __init__(self, edge_dim):
        super(TransformerObject, self).__init__()
        self.conv = TransformerConv((-1, -1), 16, edge_dim=edge_dim, beta=True, dropout=.1)

    def forward(self, x, edge_index, edge_attr):
        x = self.conv(x, edge_index, edge_attr=edge_attr)
        return x

    def reset_parameters(self):
        self.conv.reset_parameters()

class TransformerModel(nn.Module):
    def __init__(self, num_layers, edge_dim):
        super(TransformerModel, self).__init__()
        self.convlist = nn.ModuleList()
        for i in range(0, num_layers):
            conv = TransformerObject(edge_dim)
            # conv = TransformerConv((-1, -1), 16, edge_dim=edge_dim, beta=True, dropout=.1)
            self.convlist.append(conv)
        self.linear_layer = nn.Linear(16, 1)

    def forward(self, x, edge_index, edge_attr):
        for i, _ in enumerate(self.convlist):
            x = self.convlist[i](x, edge_index, edge_attr)
        x = self.linear_layer(x)
        return x

transformer_model = TransformerModel(2, 3)
transformer_model = to_hetero(transformer_model, data.metadata(), aggr='sum')

out_transformer = transformer_model(data.x_dict, data.edge_index_dict, data.edge_attr_dict)
print("finished")

The problem does not occur if I call TransformerConv(...) instead of TransformerObject(...) by simply uncommenting the line.

The error message:

Exception has occurred: ValueError
Found indices in 'edge_index' that are larger than 854 (got 963). Please ensure that all indices in 'edge_index' point to valid indices in the interval [0, 855) in your node feature matrix and try again.
IndexError: index out of range in self

Perhaps this issue is related to #6863

Environment

• PyG version: 2.2.0 (torch-geometric)
• PyTorch version: 1.13.1
• OS: MacOS
• Python version: Python 3.10.9
• CUDA/cuDNN version: no
• How you installed PyTorch and PyG (conda, pip, source): conda and pip
• Any other relevant information (e.g., version of torch-scatter): reproduced the issue on Linux with Python 3.10.0

[EdisonLeeeee]

Thank you for reporting. Currently, only MessagePassing-based layers are supported as an edge-level submodule for to_hetero. If any other modules are used (e.g., torch.nn.Module), they will be treated as node-level submodules, which may result in unexpected errors.

[rusty1s]

Yeah, this should be fixable by letting TransformerObject inherit from MessagePassing.