DataParallel imbalanced memory usage

[soulless]

I am having a similar issue. Could this be because the loss calculation is not done in the forward function?

Imbalanced GPU memory usage training LSTM

I found the reason, it is because we collect the output back to one gpu and calculate loss there. If move loss calculation into model.forward(), the problem is resolved.