Issue with dgl.mean_nodes when used with GATConv

Hi, I am trying to build a graph classifier based on GATConv layers. here is a screenshot of my classifier code. the problem is when I use dgl.mean_nodes(graph, 'h_n') I don’t get a value for each graph. on the other hand, when I use dgl.mean_nodes(graph, 'h_n') with GCNConv it works as expected. could anyone give some explanation about this behavior?

hi, what did you get in GATConv? what’s the shape of it?

(num_nodes, 1) instead of (num_graphs, 1)

I tried to repro according to your classifier and works as expected. here’s the whole demo(including dataset, train and so on): dgl_discuss_2351 · GitHub. You could run it directly. could you check the difference between yours and mine. If no more difference could be found, could you share your whole reproducible demo?

Hi @Rhett-Ying, figured it out. the problem in my training loop I used dgl.add_self_loops(batched_graph) which it returns a graph. Thanks for help