I had tried put the loss inside of forward() before. It did mitigate this imbalanced issue but not fully resolved it. BTW, I wonder if PyTorch has plans to implement the functionality to store model parameters in CPU instead of one main GPU? I’m training production scale data which is super big. If the GPU memory is used by 99% after mounting the model, there is little room for me to mount data anymore, which totally lose the scalability of multiple GPU training.