Get subset of graph nodes within batch

[manavkulshrestha]

Hello, I have a question for how to do something with batching. I have the following Data objects that my dataloader yields (batch size 1): DataBatch(x=[15, 4], edge_index=[2, 120], y=[5], batch=[15], ptr=[2]). Now, once my model gives me the output, I only want to get the loss by comparing the node predictions on the first len(y) nodes in each graph since those nodes are different and my model output is only defined for them. That is, something like

out = model(batch.x, batch.edge_index)
loss = loss_fn(out[:len(batch.y)], batch.y)

But, this of course, doesn't really work when we have batching since multiple graph instances are concatenated with each other and the out values I care about are now interleaved. Is there some way to achieve what I need? I cannot think of a good way to do this

[rusty1s]

The easiest way to achieve this is to add a mask to each data object, that holds the valid nodes you want to do for loss computation. Then, you can simply do

loss = loss_fn(out[batch.valid_mask], batch.y[batch.valid_mask])