That’s a very intersting solution to the problem. I only saw the weight parameter in the loss func, but don’t think it fits this situation. If it works, it will scale down nicely. I give a try and let you know the result. Thank you.
When nn.CrossEntropyLoss(reduction='sum') is used, the model simple does not converge with torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9).
Once swith back to mean, the same model converge within 5 epochs tested on MNIST.
It seems the sum option is only used in evaluation for computing the overall loss, and nobody use it on training.