How to do attention along different graphs?

[StefanIsSmart]

For example, I have C graphs from a mini-batch. (For example, they could be from different conv_layers.)

So the data is (L, N, C) -> L is different from different mini-batch, N is the Dimension of features, and C is the number of graphs.

What I want to do is do the attention along the C dim.

Is there any function in PyG that could do that?

[rusty1s]

You should be able to leverage torch.nn.MultiHeadAttention for this.