Pooling over heterographs

[lenianiva]

Is there a way to make the GlobalPool layer work for GNNHeteroGraphs? I only need to pool over one kind of nodes.

My current solution:

x = NNlib.scatter(model.pool.aggr, x[:kind], graph_indicator(g, :kind))

Is there a way to pool over all types of nodes or some types of nodes?