How to average the usage on multi GPU?

[JuanFMontesinos (Juan Montesinos)]

DataParallel imbalanced memory usage

Hi there @albanD, @Yuzhou_Song I noticed there is an small mistake with the code you provided: It’s necessary to unsqueeze loss inside forward pass to DataParallel were able to build loss back. Loss provided by PyTorch loss functions seems not to have dimensions, and DataParallel mount batch back in dim 0 by default. It also requires to class super class not to raise up with an error. Anyway thank you very much. Awesome help. class FullModel(nn.Module): def __init__(self, model, loss): …

Try this one

It’s roughly the same