How to use GATv2Conv with batch?

[AlbertoWang]

I wonder to know, if it is possible to use torch_geometric.nn.conv.gatv2_conv.GATv2Conv with a input of batched data?
My input x is a tensor with shape (batch_size, node_number, feature_dim), and my input edge_index is a tensor with shape (batch_size, node_number, node_number) which is a batch of adjacent matrixes. How can I use GATv2Conv for the batched data?

BTW, is there any demo doc for these models or layers implemented in pyg?

[rusty1s]

See here for a tutorial on how PyG handles mini-batches. In general, this scheme is applicable for any GNN layer integrated into PyG.

data_list = [Data(x=..., edge_index=...), Data(x=..., edge_index=...)]
loader = DataLoader(data_list, batch_size=32)

conv = GATv2Conv(...)

batch = next(iter(loader))
out = conv(batch.x, batch.edge_index)

Let me know if this resolves your issues:

[AlbertoWang]

Thanks a lot!
And I wonder what is the shape of edge_index? (edge_number, 2) or (2, edge_number)?
If each edge_index has different number of edges(data_list[0].edge_index.size() is (20, 2) and data_list[1].edge_index.size() is 30, 2 for example ), can these graphs wrap in a mini-batches?

[rusty1s]

edge_index needs to be of shape [2, num_edges], and supports varying number of nodes and edges per example in a mini-batch.

[AlbertoWang]

Thanks for ur helping!