Weight missmatch for specific batch sizes

[afonso-sousa]

Hi. I have a working model using a RGCN, like this:

class T5GNNAdapt(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.conv1 = FastRGCNConv(
            config.d_model,
            config.d_ff,
            num_relations=config.num_relations,
        )
        self.conv2 = FastRGCNConv(
            config.d_ff,
            config.d_ff,
            num_relations=config.num_relations,
        )
        self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)

        self.layer_norm = T5LayerNorm(
            config.d_model, eps=config.layer_norm_epsilon
        )
        self.dropout = nn.Dropout(config.dropout_rate)
        self.dropout_gnn = nn.Dropout(config.dropout_rate)

    def forward(self, hidden_states, graphs):
            y = F.elu(
                self.conv1(
                    graphs.x, graphs.edge_index, edge_type=graphs.edge_type
                )
            )
            y = self.dropout_gnn(y)

            y = self.conv2(
                graphs.x, graphs.edge_index, edge_type=graphs.edge_type
            )

            y = self.wo(y)

            y = global_mean_pool(y, graphs.batch)

            norm_y = self.layer_norm(y)
            norm_y = norm_y.unsqueeze(1).expand_as(hidden_states)

            layer_output = hidden_states + self.dropout(norm_y)
            return layer_output

It works fine using batch sizes of 4, 8 and 16, but when going for 32 and 64, it throws:

File "/home/admin/projects/graph_adapter/graph_t5.py", line 66, in forward
    self.conv1(graphs.x, graphs.edge_index, edge_type=graphs.edge_type)
  File "/home/admin/miniconda/envs/dp_adapter/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/admin/miniconda/envs/dp_adapter/lib/python3.9/site-packages/torch_geometric/nn/conv/rgcn_conv.py", line 311, in forward
    out = self.propagate(edge_index, x=x_l, edge_type=edge_type, size=size)
  File "/home/admin/miniconda/envs/dp_adapter/lib/python3.9/site-packages/torch_geometric/nn/conv/message_passing.py", line 467, in propagate
    out = self.message(**msg_kwargs)
  File "/home/admin/miniconda/envs/dp_adapter/lib/python3.9/site-packages/torch_geometric/nn/conv/rgcn_conv.py", line 344, in message
    return weight.view(-1, self.out_channels)[weight_index]
IndexError: index 23115 is out of bounds for dimension 0 with size 23040

I have no intuition on what might be causing the problem as it works for smaller batch sizes. Everything seems to be ok with the batch itself. Is there any batch size-dependent operation inside RGCNConv?

[rusty1s]

Thanks for reporting. Do you have a minimal script to reproduce this? This would help a lot to track this issue down.