Loss is only normalized w.r.t. batch size in train step

[greydanus]

Hey Lighting devs, thanks for a great library. I believe I have found a bug in this file:

In this line, the loss is being normalized according to the first (batch) dimension of x. However, if output_size is greater than one, then the sum over MSE errors will have run across both the batch dimension and the feature dimension (which is of length output_size.

Meanwhile, when the validation loss is computed, a reduce mean is being run across all the MSE values. This causes val_loss to differ from train_loss by a factor of output_size

[otaj]

Hi, @greydanus! We are currently going through a major revision. Please, give us some time to finalize the revision, or, sign up for the part of the revision (we're going to more than welcome every helping hand :zap:) and fix possible old bugs that you've discovered while doing that revision. The revision issue is #839 :nut_and_bolt:

[stale[bot]]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.