Gather from different graphs in batch by index

[antoni4040]

I want to pass some graphs through a model and then gather by index from each graph. I've been doing something like this:

 graphActionDataset = graphDataset(graphsForAction)
out = [action(j).flatten() for j in graphActionDataset]
out = torch.nn.utils.rnn.pad_sequence(out, batch_first=True, padding_value=0.0)
outValues = torch.gather(input=out, dim=1, index=transitionsForAction.unsqueeze(1))

Which is annoying as I need to create a dataset and then pad the tensors as torch.gather wants equal sized tensors and the graphs have different sizes. Is there an easier way to do this, preferably from a Batch?

[rusty1s]

How about using torch_geometric.utils.to_dense_batch for this?

[antoni4040]

Yeap, thanks!