Is loss averaged correctly?

[mallamanis]

I suspect (but I am not sure) that the loss might not be normalized correctly here.

MSELoss seems to be averaging the loss across the batch size. Then you divide the sum of the epoch loss by the number of samples (and not the number of minibatches).

If you have time, please check this (but I might be missing something from the code). If this is indeed wrong, the reported training/evaluation loss will be a bit off.