Assertion Error in GNN explainer

[amoschoomy]

I have 100s of graphs to be analysed, but this error only occurs on specific graphs, so I'm at a loss on where to find the issue and fix it.

File "/scratch/project_mnt/S0163/GNN/explainer.py", line 16, in generate_explanation
    explanation = explainer(data.x, data.edge_index, batch=data.batch)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch_geometric/explain/explainer.py", line 196, in __call__
    prediction = self.get_prediction(x, edge_index, **kwargs)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch_geometric/explain/explainer.py", line 115, in get_prediction
    out = self.model(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/GNN/train.py", line 79, in forward
    x = F.leaky_relu(self.conv1(x, edge_index))
                     ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch_geometric/nn/conv/gatv2_conv.py", line 302, in forward
    out = self.propagate(edge_index, x=(x_l, x_r), alpha=alpha)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch_geometric/nn/conv/message_passing.py", line 563, in propagate
    out = self.explain_message(out, **explain_msg_kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/scratch/project_mnt/S0163/gnn_env/lib/python3.11/site-packages/torch_geometric/nn/conv/message_passing.py", line 799, in explain_message
    assert inputs.size(self.node_dim) == edge_mask.size(0)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
AssertionError

This is my network:

class GATv2(torch.nn.Module):
    """
    Graph Attention Network version 2 (GATv2) module.

    Args:
        nfeat (int): Number of input features.
        nhid (int): Number of hidden units.
        nclass (int): Number of output classes.
        dropout (float): Dropout rate.
        temperature (float, optional): Temperature parameter to scale the confidence of output probabilities. Defaults to 1.1.

    Attributes:
        conv1 (GATv2Conv): First GATv2 convolutional layer.
        conv2 (GATv2Conv): Second GATv2 convolutional layer.
        classifier (torch.nn.Linear): Linear layer for classification.
        dropout (float): Dropout rate.
        temperature (float): Temperature parameter.

    Methods:
        forward(x, edge_index, batch): Forward pass of the GATv2 module.


        Forward pass of the GATv2 module.

        Args:
            x (torch.Tensor): Input features.
            edge_index (torch.Tensor): Graph edge indices.
            batch (torch.Tensor): Batch indices.

        Returns:
            torch.Tensor: Output tensor after the forward pass.

    """

    def __init__(
        self,
        nfeat: int,
        nhid: int,
        nclass: int,
        dropout: float,
        temperature: float = 1.034,
    ) -> None:
        super(GATv2, self).__init__()
        self.conv1 = GATv2Conv(nfeat, nhid, heads=32, dropout=dropout, concat=False)
        self.conv2 = GATv2Conv(nhid, nhid, heads=32, dropout=dropout, concat=False)
        self.classifier = torch.nn.Linear(nhid, nclass)
        self.dropout = dropout
        self.temperature = temperature  # scale confidence of output prob

    def forward(
        self, x: torch.Tensor, edge_index: torch.Tensor, batch: torch.Tensor
    ) -> torch.Tensor:
        x = F.leaky_relu(self.conv1(x, edge_index))
        x = F.leaky_relu(self.conv2(x, edge_index))
        x = F.dropout(x, p=self.dropout, training=self.training)  # dropout layer
        x = global_mean_pool(x, batch)  # Global pooling
        x = self.classifier(x)
        x /= self.temperature
        return x