Loss broadcasting fix

[laserkelvin]

This PR attempts to correct broadcasting issues due to shape mismatches in loss calculations.

This was brought on by the realization that broadcasting works a little differently (from some time ago) when tensor shapes are mismatched. In particular, labels come out of the pipeline with an extra dimension (e.g. [N, 1]) compared to graph readouts. The resulting behavior is actually very different from the intention:

>>> y
tensor([0.0551, 0.2665, 0.4638, 0.3288, 0.1201, 0.1515, 0.9187, 0.4527])
>>> x
tensor([[0.0961],
        [0.6194],
        [0.3628],
        [0.5289],
        [0.2046],
        [0.0989],
        [0.6621],
        [0.8717]])
>>> from torch.nn import MSELoss
>>> MSELoss()(y, x)
/python3.12/site-packages/torch/nn/modules/loss.py:610: UserWarning: Using a target size (torch.Size([8, 1])) that is different to the input size (torch.Size([8])). 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)
tensor(0.1454)
>>> MSELoss()(y.view(-1, 1), x)
tensor(0.0535)

The code changes make it so that _compute_losses methods will check if model output and label shapes are mismatched, and if they are, attempt to reshape the model outputs to match the labels' shape before computing the loss.