Problem with mask shape while using remove_self_loops

[LaurentBerder]

Hi,

I'm working with the following setup:

• platform: linux
• python: 3.7
• torch_geometric_temporal: 0.54.0
• torch_geometric: 2.2.0

I have been preparing a dataset for a while, now. It's a DynamicGraphTemporalSignal because the values of my edge features varies over time, and the idea is to build a model that will be able to forecast 30 timesteps in advance.

Here are some values to explain the shape of my dataset:

nb_nodes = 715
nb_edges = 13339
nb_node_features = 3
nb_edge_features = 2
nb_timesteps = 385
days_in = 30
days_out = 30

print(type(dataset))

<class 'torch_geometric_temporal.signal.dynamic_graph_temporal_signal.DynamicGraphTemporalSignal'>

dataset |--node_features: (nb_timesteps, nb_nodes, nb_node_features, days_in) |--node_targets: (nb_timesteps, nb_nodes, days_out) |--edge_index: (nb_timesteps, 2, nb_edges, days_in) |--edge_features: (nb_timesteps, nb_edge_features, nb_edges, days_in)

For training purposes, I made sure to split my dataset into train and test

train_dataset, test_dataset = temporal_signal_split(dataset, train_ratio=0.8)

Now my problem is the modeling. I am trying a GConvLSTM layer with the following:

from torch_geometric_temporal.nn import GConvLSTM

class DynGraphModel(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, K):
        super(DynGraphModel, self).__init__()
        self.conv = GConvLSTM(input_dim, hidden_dim, K)
        self.linear = torch.nn.Linear(hidden_dim, output_dim)
    def forward(self, x, edge_index, edge_attr, H, C):
        h_0, c_0 = self.conv(x, edge_index, edge_attr, H, C)
        h = F.relu(h_0)
        out = self.linear(h)
        return h, h_0, c_0

hidden_channels = 32

model = DynGraphModel(input_dim=nb_node_features, hidden_dim=hidden_channels, output_dim=1, K=1)
snapshot = next(iter(train_dataset))
H = None
C = None

y_hat = model(x=snapshot.x, edge_index=snapshot.edge_index, edge_attr=snapshot.edge_attr, H=H, C=C)

Which gives me the following stacktrace:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "<stdin>", line 7, in forward
  File "python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "python3.7/site-packages/torch_geometric_temporal/nn/recurrent/gconv_lstm.py", line 231, in forward
    I = self._calculate_input_gate(X, edge_index, edge_weight, H, C, lambda_max)
  File "site-packages/torch_geometric_temporal/nn/recurrent/gconv_lstm.py", line 167, in _calculate_input_gate
    I = self.conv_x_i(X, edge_index, edge_weight, lambda_max=lambda_max)
  File "python3.7/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "python3.7/site-packages/torch_geometric/nn/conv/cheb_conv.py", line 165, in forward
    batch=batch,
  File "python3.7/site-packages/torch_geometric/nn/conv/cheb_conv.py", line 122, in __norm__
    edge_index, edge_weight = remove_self_loops(edge_index, edge_weight)
  File "python3.7/site-packages/torch_geometric/utils/loop.py", line 66, in remove_self_loops
    return edge_index, edge_attr[mask]
IndexError: The shape of the mask [13339, 30] at index 0 does not match the shape of the indexed tensor [2, 13339, 30] at index 0

I understand that the issue seems points to torch_geometric rather than torch_geometric_temporal, but is that something someone already encountered?

I see no reason why the mask should be the wrong shape.

[LaurentBerder]

What's particularly frustrating is that this happens in the process of eliminating self loops, while my dataset doesn't have any:

snapshot.has_self_loops()
>>>False