nn.NNConv's edge_feature mapping has inconsistent shape

[kaifhm]

I have the following code:

edge_aggr_nn = torch.nn.Linear(edge_embed_dims, atom_embed_dims*17)
nnconv1 = NNConv(in_channels=atom_embed_dims, out_channels=17, nn=edge_aggr_nn)
nnconv1(smiles_graph.x, smiles_graph.edge_index, smiles_graph.edge_attr)

This code should run, but I get RuntimeError: The size of tensor a (4900) must match the size of tensor b (2450) at non-singleton dimension 0. According to the documentation nn must have input shape [-1, num_edge_features] and output shape [-1, in_channels * out_channels]. Notice that tensor a is twice the size of tensor b. So I increased out_feature of nn to 2*out_channels (i.e. 2*17) and it worked. My question is, why did it work? What might possibly be going on and what am I misunderstanding about the documentation?

This is with torch==2.8.0+cu126, torch-geometric==2.7.0