Don’t ignore this warning as it’s real:
UserWarning: Using a target size (torch.Size([10])) that is different to the input size (torch.Size([10, 1])). This will likely lead to incorrect results due to broadcasting. Please ensure they have the same size.
return F.mse_loss(input, target, reduction=self.reduction)
Fix the shape of the model output and the target and make sure they are matching. Otherwise the loss calculation will broadcast one tensor and will thus calculate a wrong loss value.