Backwards incompatible change to MSE for pixelwise regression

[adamjstewart]

🐛 Bug

#1937 introduces an unintended consequence: pixelwise regression is no longer supported.

To Reproduce

Run the following script:

import torch
import torchmetrics

B = 4
H = W = 3

x = torch.rand(B, H, W)
y = torch.rand(B, H, W)

torchmetrics.functional.mean_squared_error(x, y)

This results in the following error msg:

Traceback (most recent call last):
  File "test.py", line 10, in <module>
    torchmetrics.functional.mean_squared_error(x, y, num_outputs=H * W)
  File "lib/python3.10/site-packages/torchmetrics/functional/regression/mse.py", line 84, in mean_squared_error
    sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=num_outputs)
  File "lib/python3.10/site-packages/torchmetrics/functional/regression/mse.py", line 35, in _mean_squared_error_update
    _check_data_shape_to_num_outputs(preds, target, num_outputs, allow_1d_reshape=True)
  File "lib/python3.10/site-packages/torchmetrics/functional/regression/utils.py", line 31, in _check_data_shape_to_num_outputs
    raise ValueError(
ValueError: Expected both predictions and target to be either 1- or 2-dimensional tensors, but got 3 and 3.

Expected behavior

I would expect the MSE metrics to support pixelwise regression (predicting a single regression value for each pixel in an image). The above script works fine with torchmetrics 1.0.3.

Environment

• TorchMetrics version: 1.1.0, spack
• Python & PyTorch Version: 3.10.10, 2.1.0
• Any other relevant information such as OS: macOS

Additional context

@SkafteNicki @Borda @justusschock