Break batch embeddings into each sample

[zcaicaros]

Hi,
How can I break the batch node embeddings? After processing a data.Batch batch of data, the node embedding has dimension [batch_size * #nodes_in_each_graph, output_dim], but how can I break it into [batch_size, nodes_#in_each_graph, output_dim]?

[rusty1s]

You can use torch_geometric.utils.to_dense_batch to achieve that. Note that for graphs with varying number of nodes, it will return a tensor of shape [batch_size, max_num_nodes, output_dim], filled with zeros for smaller graphs.