how to train the timesteps dataset by edge_net?

[scalaboy]

as show in this example code :

import torch
from torch import Tensor
from torch.nn import Sequential, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
def init(self, in_channels, out_channels):
super().init(aggr="max") # "Max" aggregation.
self.mlp = Sequential(
Linear(2 * in_channels, out_channels),
ReLU(),
Linear(out_channels, out_channels),
)

def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
    # x: Node feature matrix of shape [num_nodes, in_channels]
    # edge_index: Graph connectivity matrix of shape [2, num_edges]
    return self.propagate(edge_index, x=x)  # shape [num_nodes, out_channels]

def message(self, x_j: Tensor, x_i: Tensor) -> Tensor:
    # x_j: Source node features of shape [num_edges, in_channels]
    # x_i: Target node features of shape [num_edges, in_channels]
    edge_features = torch.cat([x_i, x_j - x_i], dim=-1)
    return self.mlp(edge_features)  # shape [num_edges, out_channels]

how ever,the shape x: Node feature matrix of shape [num_nodes, in_channels]
where my x shape is : Node feature matrix of shape [num_nodes,timesteps, in_channels],

[rusty1s]

Please see my answer here: #4084