Temporal GNN not working in WindmillOutputMediumDatasetLoader dataset

[osllogon]

I have been trying to train a temporal GNN with the following structure on WindmillOutputMediumDatasetLoader dataset:

`class RecurrentGCN(torch.nn.Module):

def __init__(self, node_features: int) -> None:
    # call super class constructor
    super().__init__()

    # define layers
    self.recurrent: torch.nn.Module = GConvLSTM(node_features, 32, 2)
    self.relu: torch.nn.Module = torch.nn.ReLU()
    self.linear: torch.nn.Module = torch.nn.Linear(32, 1)

def forward(self, x: torch.Tensor, edge_index: torch.Tensor, edge_weight: torch.Tensor) -> torch.Tensor:
    # get outputs
    outputs = self.recurrent(x, edge_index, edge_weight)[0]
    outputs = self.relu(outputs)
    outputs = self.linear(outputs)

    return outputs`

But I find impossible to optimize. I have tried many different learning rates and optimizers. Can anyone help me with this or give me any tips? I have tried the same with other datasets from the library and I found no problem

[osllogon]

Is there any solution or update for this?