Another issue with assembling the batch in the forward function is that you end up doing the same work multiple times, depending on the number of GPUs you are using (the forward function is called N times).



 alsuhr:

I think I know what the issue is – does DataParallel require that the input to the forward calls be tensors so it can distribute them?


Yes. If you pass the input batch (as a tensor) to nn.DataParallel, it will split along the batch dimension and distribute the smaller batches to participating GPUs.