We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Backwards incompatible change to MSE for pixelwise regression
## 🐛 Bug

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

### To Reproduce

Run the following script:
```python
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 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `src/torchmetrics/functional/regression/mse.py`
Content:
```
1 # Copyright The Lightning team.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 from typing import Tuple, Union
15 
16 import torch
17 from torch import Tensor
18 
19 from torchmetrics.functional.regression.utils import _check_data_shape_to_num_outputs
20 from torchmetrics.utilities.checks import _check_same_shape
21 
22 
23 def _mean_squared_error_update(preds: Tensor, target: Tensor, num_outputs: int) -> Tuple[Tensor, int]:
24     """Update and returns variables required to compute Mean Squared Error.
25 
26     Check for same shape of input tensors.
27 
28     Args:
29         preds: Predicted tensor
30         target: Ground truth tensor
31         num_outputs: Number of outputs in multioutput setting
32 
33     """
34     _check_same_shape(preds, target)
35     _check_data_shape_to_num_outputs(preds, target, num_outputs, allow_1d_reshape=True)
36     if num_outputs == 1:
37         preds = preds.view(-1)
38         target = target.view(-1)
39     diff = preds - target
40     sum_squared_error = torch.sum(diff * diff, dim=0)
41     n_obs = target.shape[0]
42     return sum_squared_error, n_obs
43 
44 
45 def _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: Union[int, Tensor], squared: bool = True) -> Tensor:
46     """Compute Mean Squared Error.
47 
48     Args:
49         sum_squared_error: Sum of square of errors over all observations
50         n_obs: Number of predictions or observations
51         squared: Returns RMSE value if set to False.
52 
53     Example:
54         >>> preds = torch.tensor([0., 1, 2, 3])
55         >>> target = torch.tensor([0., 1, 2, 2])
56         >>> sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=1)
57         >>> _mean_squared_error_compute(sum_squared_error, n_obs)
58         tensor(0.2500)
59 
60     """
61     return sum_squared_error / n_obs if squared else torch.sqrt(sum_squared_error / n_obs)
62 
63 
64 def mean_squared_error(preds: Tensor, target: Tensor, squared: bool = True, num_outputs: int = 1) -> Tensor:
65     """Compute mean squared error.
66 
67     Args:
68         preds: estimated labels
69         target: ground truth labels
70         squared: returns RMSE value if set to False
71         num_outputs: Number of outputs in multioutput setting
72 
73     Return:
74         Tensor with MSE
75 
76     Example:
77         >>> from torchmetrics.functional.regression import mean_squared_error
78         >>> x = torch.tensor([0., 1, 2, 3])
79         >>> y = torch.tensor([0., 1, 2, 2])
80         >>> mean_squared_error(x, y)
81         tensor(0.2500)
82 
83     """
84     sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=num_outputs)
85     return _mean_squared_error_compute(sum_squared_error, n_obs, squared=squared)
86
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/src/torchmetrics/functional/regression/mse.py b/src/torchmetrics/functional/regression/mse.py
--- a/src/torchmetrics/functional/regression/mse.py
+++ b/src/torchmetrics/functional/regression/mse.py
@@ -16,7 +16,6 @@
 import torch
 from torch import Tensor
 
-from torchmetrics.functional.regression.utils import _check_data_shape_to_num_outputs
 from torchmetrics.utilities.checks import _check_same_shape
 
 
@@ -32,7 +31,6 @@
 
     """
     _check_same_shape(preds, target)
-    _check_data_shape_to_num_outputs(preds, target, num_outputs, allow_1d_reshape=True)
     if num_outputs == 1:
         preds = preds.view(-1)
         target = target.view(-1)

{"golden_diff": "diff --git a/src/torchmetrics/functional/regression/mse.py b/src/torchmetrics/functional/regression/mse.py\n--- a/src/torchmetrics/functional/regression/mse.py\n+++ b/src/torchmetrics/functional/regression/mse.py\n@@ -16,7 +16,6 @@\n import torch\n from torch import Tensor\n \n-from torchmetrics.functional.regression.utils import _check_data_shape_to_num_outputs\n from torchmetrics.utilities.checks import _check_same_shape\n \n \n@@ -32,7 +31,6 @@\n \n     \"\"\"\n     _check_same_shape(preds, target)\n-    _check_data_shape_to_num_outputs(preds, target, num_outputs, allow_1d_reshape=True)\n     if num_outputs == 1:\n         preds = preds.view(-1)\n         target = target.view(-1)\n", "issue": "Backwards incompatible change to MSE for pixelwise regression\n## \ud83d\udc1b Bug\r\n\r\n#1937 introduces an unintended consequence: pixelwise regression is no longer supported.\r\n\r\n### To Reproduce\r\n\r\nRun the following script:\r\n```python\r\nimport torch\r\nimport torchmetrics\r\n\r\nB = 4\r\nH = W = 3\r\n\r\nx = torch.rand(B, H, W)\r\ny = torch.rand(B, H, W)\r\n\r\ntorchmetrics.functional.mean_squared_error(x, y)\r\n```\r\nThis results in the following error msg:\r\n```\r\nTraceback (most recent call last):\r\n  File \"test.py\", line 10, in <module>\r\n    torchmetrics.functional.mean_squared_error(x, y, num_outputs=H * W)\r\n  File \"lib/python3.10/site-packages/torchmetrics/functional/regression/mse.py\", line 84, in mean_squared_error\r\n    sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=num_outputs)\r\n  File \"lib/python3.10/site-packages/torchmetrics/functional/regression/mse.py\", line 35, in _mean_squared_error_update\r\n    _check_data_shape_to_num_outputs(preds, target, num_outputs, allow_1d_reshape=True)\r\n  File \"lib/python3.10/site-packages/torchmetrics/functional/regression/utils.py\", line 31, in _check_data_shape_to_num_outputs\r\n    raise ValueError(\r\nValueError: Expected both predictions and target to be either 1- or 2-dimensional tensors, but got 3 and 3.\r\n```\r\n\r\n### Expected behavior\r\n\r\nI 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.\r\n\r\n### Environment\r\n\r\n- TorchMetrics version: 1.1.0, spack\r\n- Python & PyTorch Version: 3.10.10, 2.1.0\r\n- Any other relevant information such as OS: macOS\r\n\r\n### Additional context\r\n\r\n@SkafteNicki @Borda @justusschock \n", "before_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.functional.regression.utils import _check_data_shape_to_num_outputs\nfrom torchmetrics.utilities.checks import _check_same_shape\n\n\ndef _mean_squared_error_update(preds: Tensor, target: Tensor, num_outputs: int) -> Tuple[Tensor, int]:\n    \"\"\"Update and returns variables required to compute Mean Squared Error.\n\n    Check for same shape of input tensors.\n\n    Args:\n        preds: Predicted tensor\n        target: Ground truth tensor\n        num_outputs: Number of outputs in multioutput setting\n\n    \"\"\"\n    _check_same_shape(preds, target)\n    _check_data_shape_to_num_outputs(preds, target, num_outputs, allow_1d_reshape=True)\n    if num_outputs == 1:\n        preds = preds.view(-1)\n        target = target.view(-1)\n    diff = preds - target\n    sum_squared_error = torch.sum(diff * diff, dim=0)\n    n_obs = target.shape[0]\n    return sum_squared_error, n_obs\n\n\ndef _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: Union[int, Tensor], squared: bool = True) -> Tensor:\n    \"\"\"Compute Mean Squared Error.\n\n    Args:\n        sum_squared_error: Sum of square of errors over all observations\n        n_obs: Number of predictions or observations\n        squared: Returns RMSE value if set to False.\n\n    Example:\n        >>> preds = torch.tensor([0., 1, 2, 3])\n        >>> target = torch.tensor([0., 1, 2, 2])\n        >>> sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=1)\n        >>> _mean_squared_error_compute(sum_squared_error, n_obs)\n        tensor(0.2500)\n\n    \"\"\"\n    return sum_squared_error / n_obs if squared else torch.sqrt(sum_squared_error / n_obs)\n\n\ndef mean_squared_error(preds: Tensor, target: Tensor, squared: bool = True, num_outputs: int = 1) -> Tensor:\n    \"\"\"Compute mean squared error.\n\n    Args:\n        preds: estimated labels\n        target: ground truth labels\n        squared: returns RMSE value if set to False\n        num_outputs: Number of outputs in multioutput setting\n\n    Return:\n        Tensor with MSE\n\n    Example:\n        >>> from torchmetrics.functional.regression import mean_squared_error\n        >>> x = torch.tensor([0., 1, 2, 3])\n        >>> y = torch.tensor([0., 1, 2, 2])\n        >>> mean_squared_error(x, y)\n        tensor(0.2500)\n\n    \"\"\"\n    sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=num_outputs)\n    return _mean_squared_error_compute(sum_squared_error, n_obs, squared=squared)\n", "path": "src/torchmetrics/functional/regression/mse.py"}], "after_files": [{"content": "# Copyright The Lightning team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torchmetrics.utilities.checks import _check_same_shape\n\n\ndef _mean_squared_error_update(preds: Tensor, target: Tensor, num_outputs: int) -> Tuple[Tensor, int]:\n    \"\"\"Update and returns variables required to compute Mean Squared Error.\n\n    Check for same shape of input tensors.\n\n    Args:\n        preds: Predicted tensor\n        target: Ground truth tensor\n        num_outputs: Number of outputs in multioutput setting\n\n    \"\"\"\n    _check_same_shape(preds, target)\n    if num_outputs == 1:\n        preds = preds.view(-1)\n        target = target.view(-1)\n    diff = preds - target\n    sum_squared_error = torch.sum(diff * diff, dim=0)\n    n_obs = target.shape[0]\n    return sum_squared_error, n_obs\n\n\ndef _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: Union[int, Tensor], squared: bool = True) -> Tensor:\n    \"\"\"Compute Mean Squared Error.\n\n    Args:\n        sum_squared_error: Sum of square of errors over all observations\n        n_obs: Number of predictions or observations\n        squared: Returns RMSE value if set to False.\n\n    Example:\n        >>> preds = torch.tensor([0., 1, 2, 3])\n        >>> target = torch.tensor([0., 1, 2, 2])\n        >>> sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=1)\n        >>> _mean_squared_error_compute(sum_squared_error, n_obs)\n        tensor(0.2500)\n\n    \"\"\"\n    return sum_squared_error / n_obs if squared else torch.sqrt(sum_squared_error / n_obs)\n\n\ndef mean_squared_error(preds: Tensor, target: Tensor, squared: bool = True, num_outputs: int = 1) -> Tensor:\n    \"\"\"Compute mean squared error.\n\n    Args:\n        preds: estimated labels\n        target: ground truth labels\n        squared: returns RMSE value if set to False\n        num_outputs: Number of outputs in multioutput setting\n\n    Return:\n        Tensor with MSE\n\n    Example:\n        >>> from torchmetrics.functional.regression import mean_squared_error\n        >>> x = torch.tensor([0., 1, 2, 3])\n        >>> y = torch.tensor([0., 1, 2, 2])\n        >>> mean_squared_error(x, y)\n        tensor(0.2500)\n\n    \"\"\"\n    sum_squared_error, n_obs = _mean_squared_error_update(preds, target, num_outputs=num_outputs)\n    return _mean_squared_error_compute(sum_squared_error, n_obs, squared=squared)\n", "path": "src/torchmetrics/functional/regression/mse.py"}]}