CompRhys:

Naively using nn.DataParallel(model) gives me an error as I believe it requires the sub-tensors it splits up the tensor into to all be the same size which is not the case here as while each GPU gets 64 graphs the number of nodes varies.


I am inclined to think this shouldn’t cause a problem for DataParallel. Because DataParallel would replicate your model and scatter your input on the first (batch) dimension. Even though different input element (graph) can contain different node, as long as your model forward function can handle it properly, it should not hit error.
Could you please share a min repro for the error you saw?