If you are using nn.DataParallel the model will be replicated to each GPU and each model will get a chunk of your input batch.
The output will be gathered on the default device, so most likely you wouldn’t have to change anything.
However, I’m not sure about the use case.
How are you calculating the memory consumption and is this operation differentiable?
I assume it’s not differentiable so that your accumulated loss will in fact just be the nn.CrossEntropyLoss.