Expected input batch_size (24) to match target batch_size (18).

[KOarch]

My error is at the criterion(out, one_hot_labels) because in the data:
GCNModel out.shape: torch.Size([24, 6])
train_model out.shape: torch.Size([24, 6])
train_model data.y.shape: torch.Size([18, 2])
train_model data.edge_attr.shape: torch.Size([18, 3])
The tensor of out.shape is the output obtained by data.x.shape
My data.y.shape is obtained from data.edge_attr

##############
out = model(data.x, data.edge_index, data.edge_attr)
one_hot_labels = convert_to_one_hot(data.y, output_dim)

print("train_model out.shape:", out.shape)
print("train_model data.y.shape:", data.y.shape)
print("train_model data.edge_attr.shape:", data.edge_attr.shape)
print("train_model data.edge_attr:", data.edge_attr)
print("train_model one_hot_labels.shape:", one_hot_labels.shape)

loss = criterion(out, one_hot_labels)

What can I do to fix the mismatch between the number of out and one_hot_labels
Am I modifying the model or am I modifying the dimensions of the output???
################## model #################
class GCNModel(nn.Module):
def init(self, in_channels, hidden_dim, out_channels, edge_dim):
super(GCNModel, self).init()
self.conv1 = Edge_GCNConv(in_channels=in_channels, out_channels=hidden_dim, edge_dim=edge_dim)
self.conv2 = Edge_GCNConv(in_channels=hidden_dim, out_channels=out_channels, edge_dim=edge_dim)
self.batch_norm1 = nn.BatchNorm1d(hidden_dim) # 添加批归一化层
self.batch_norm2 = nn.BatchNorm1d(out_channels) # 添加批归一化层
self.linear = nn.Linear(out_channels, out_channels) # output_dim是标签的维度

def forward(self, x, edge_index, edge_attr):
    x, edge_index, edge_attr = x, edge_index, edge_attr
    
    x1 = self.conv1(x, edge_index, edge_attr)
    
    x1 = self.batch_norm1(x1)
    x1 = F.relu(x1)
    x1 = F.dropout(x1, p=0.1, training=self.training)
    # print("GCNModel x1:", x1)
    # print("GCNModel x1.shape:", x1.shape)
    

    x2 = self.conv2(x1, edge_index, edge_attr)
    x2 = self.batch_norm2(x2)
    x2 = F.relu(x2)
    
    # print("GCNModel x2:", x2)
    # print("GCNModel x2.shape:", x2.shape)
    
    x2 = F.dropout(x2, p=0.1, training=self.training)
    
    out = self.linear(x2)
    print("GCNModel out:", out)
    print("GCNModel out.shape:", out.shape)
    return out

The out of this model is calculated from the node features,
data.x.shape:torch.Size([24, 6])
GCNModel out.shape: torch.Size([24, 6])

[rusty1s]

You need to map out to edge-level space, e.g., via

self.mlp(torch.cat([x[edge_index[0], edge_index[1], edge_attr], dim=1))