prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Add option to compute root_mean_squared_error
  ## 🚀 Feature
  Allow the user to choose between MSE and RMSE.
  
  ### Motivation
  In a physical domain the RMSE, which is essentially the mean of distances, may be significantly more intuitive than the MSE. Therefore, it would be nice to have the option to choose the preferd metric.
  
  ### Pitch
  Similar to the implementation in [scikit-learn](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html#sklearn.metrics.mean_squared_error) one could simply pass `squared=False` to the `MeanSquaredError` module or the `mean_squared_error` function.
  
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `torchmetrics/functional/regression/mean_squared_error.py`
  Content:
  ```
  1 # Copyright The PyTorch 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
  15 
  16 import torch
  17 from torch import Tensor
  18 
  19 from torchmetrics.utilities.checks import _check_same_shape
  20 
  21 
  22 def _mean_squared_error_update(preds: Tensor, target: Tensor) -> Tuple[Tensor, int]:
  23     _check_same_shape(preds, target)
  24     diff = preds - target
  25     sum_squared_error = torch.sum(diff * diff)
  26     n_obs = target.numel()
  27     return sum_squared_error, n_obs
  28 
  29 
  30 def _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int) -> Tensor:
  31     return sum_squared_error / n_obs
  32 
  33 
  34 def mean_squared_error(preds: Tensor, target: Tensor) -> Tensor:
  35     """
  36     Computes mean squared error
  37 
  38     Args:
  39         preds: estimated labels
  40         target: ground truth labels
  41 
  42     Return:
  43         Tensor with MSE
  44 
  45     Example:
  46         >>> from torchmetrics.functional import mean_squared_error
  47         >>> x = torch.tensor([0., 1, 2, 3])
  48         >>> y = torch.tensor([0., 1, 2, 2])
  49         >>> mean_squared_error(x, y)
  50         tensor(0.2500)
  51     """
  52     sum_squared_error, n_obs = _mean_squared_error_update(preds, target)
  53     return _mean_squared_error_compute(sum_squared_error, n_obs)
  54
  ```
  
  Path: `torchmetrics/regression/mean_squared_error.py`
  Content:
  ```
  1 # Copyright The PyTorch 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 Any, Callable, Optional
  15 
  16 import torch
  17 from torch import Tensor, tensor
  18 
  19 from torchmetrics.functional.regression.mean_squared_error import (
  20     _mean_squared_error_compute,
  21     _mean_squared_error_update,
  22 )
  23 from torchmetrics.metric import Metric
  24 
  25 
  26 class MeanSquaredError(Metric):
  27     r"""
  28     Computes `mean squared error <https://en.wikipedia.org/wiki/Mean_squared_error>`_ (MSE):
  29 
  30     .. math:: \text{MSE} = \frac{1}{N}\sum_i^N(y_i - \hat{y_i})^2
  31 
  32     Where :math:`y` is a tensor of target values, and :math:`\hat{y}` is a tensor of predictions.
  33 
  34     Args:
  35         compute_on_step:
  36             Forward only calls ``update()`` and return None if this is set to False. default: True
  37         dist_sync_on_step:
  38             Synchronize metric state across processes at each ``forward()``
  39             before returning the value at the step. default: False
  40         process_group:
  41             Specify the process group on which synchronization is called. default: None (which selects the entire world)
  42 
  43     Example:
  44         >>> from torchmetrics import MeanSquaredError
  45         >>> target = torch.tensor([2.5, 5.0, 4.0, 8.0])
  46         >>> preds = torch.tensor([3.0, 5.0, 2.5, 7.0])
  47         >>> mean_squared_error = MeanSquaredError()
  48         >>> mean_squared_error(preds, target)
  49         tensor(0.8750)
  50 
  51     """
  52 
  53     def __init__(
  54         self,
  55         compute_on_step: bool = True,
  56         dist_sync_on_step: bool = False,
  57         process_group: Optional[Any] = None,
  58         dist_sync_fn: Callable = None,
  59     ):
  60         super().__init__(
  61             compute_on_step=compute_on_step,
  62             dist_sync_on_step=dist_sync_on_step,
  63             process_group=process_group,
  64             dist_sync_fn=dist_sync_fn,
  65         )
  66 
  67         self.add_state("sum_squared_error", default=tensor(0.0), dist_reduce_fx="sum")
  68         self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
  69 
  70     def update(self, preds: Tensor, target: Tensor):
  71         """
  72         Update state with predictions and targets.
  73 
  74         Args:
  75             preds: Predictions from model
  76             target: Ground truth values
  77         """
  78         sum_squared_error, n_obs = _mean_squared_error_update(preds, target)
  79 
  80         self.sum_squared_error += sum_squared_error
  81         self.total += n_obs
  82 
  83     def compute(self):
  84         """
  85         Computes mean squared error over state.
  86         """
  87         return _mean_squared_error_compute(self.sum_squared_error, self.total)
  88 
  89     @property
  90     def is_differentiable(self):
  91         return True
  92
  ```
  --- 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()
  
  ```
  
golden_diff: diff --git a/torchmetrics/functional/regression/mean_squared_error.py b/torchmetrics/functional/regression/mean_squared_error.py
  --- a/torchmetrics/functional/regression/mean_squared_error.py
  +++ b/torchmetrics/functional/regression/mean_squared_error.py
  @@ -27,17 +27,18 @@
       return sum_squared_error, n_obs
   
   
  -def _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int) -> Tensor:
  -    return sum_squared_error / n_obs
  +def _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int, squared: bool = True) -> Tensor:
  +    return sum_squared_error / n_obs if squared else torch.sqrt(sum_squared_error / n_obs)
   
   
  -def mean_squared_error(preds: Tensor, target: Tensor) -> Tensor:
  +def mean_squared_error(preds: Tensor, target: Tensor, squared: bool = True) -> Tensor:
       """
       Computes mean squared error
   
       Args:
           preds: estimated labels
           target: ground truth labels
  +        squared: returns RMSE value if set to False
   
       Return:
           Tensor with MSE
  @@ -50,4 +51,4 @@
           tensor(0.2500)
       """
       sum_squared_error, n_obs = _mean_squared_error_update(preds, target)
  -    return _mean_squared_error_compute(sum_squared_error, n_obs)
  +    return _mean_squared_error_compute(sum_squared_error, n_obs, squared=squared)
  diff --git a/torchmetrics/regression/mean_squared_error.py b/torchmetrics/regression/mean_squared_error.py
  --- a/torchmetrics/regression/mean_squared_error.py
  +++ b/torchmetrics/regression/mean_squared_error.py
  @@ -39,6 +39,8 @@
               before returning the value at the step. default: False
           process_group:
               Specify the process group on which synchronization is called. default: None (which selects the entire world)
  +        squared:
  +            If True returns MSE value, if False returns RMSE value.
   
       Example:
           >>> from torchmetrics import MeanSquaredError
  @@ -56,6 +58,7 @@
           dist_sync_on_step: bool = False,
           process_group: Optional[Any] = None,
           dist_sync_fn: Callable = None,
  +        squared: bool = True,
       ):
           super().__init__(
               compute_on_step=compute_on_step,
  @@ -66,6 +69,7 @@
   
           self.add_state("sum_squared_error", default=tensor(0.0), dist_reduce_fx="sum")
           self.add_state("total", default=tensor(0), dist_reduce_fx="sum")
  +        self.squared = squared
   
       def update(self, preds: Tensor, target: Tensor):
           """
  @@ -84,7 +88,7 @@
           """
           Computes mean squared error over state.
           """
  -        return _mean_squared_error_compute(self.sum_squared_error, self.total)
  +        return _mean_squared_error_compute(self.sum_squared_error, self.total, squared=self.squared)
   
       @property
       def is_differentiable(self):
  
verification_info: {"golden_diff": "diff --git a/torchmetrics/functional/regression/mean_squared_error.py b/torchmetrics/functional/regression/mean_squared_error.py\n--- a/torchmetrics/functional/regression/mean_squared_error.py\n+++ b/torchmetrics/functional/regression/mean_squared_error.py\n@@ -27,17 +27,18 @@\n     return sum_squared_error, n_obs\n \n \n-def _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int) -> Tensor:\n-    return sum_squared_error / n_obs\n+def _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int, squared: bool = True) -> Tensor:\n+    return sum_squared_error / n_obs if squared else torch.sqrt(sum_squared_error / n_obs)\n \n \n-def mean_squared_error(preds: Tensor, target: Tensor) -> Tensor:\n+def mean_squared_error(preds: Tensor, target: Tensor, squared: bool = True) -> Tensor:\n     \"\"\"\n     Computes 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 \n     Return:\n         Tensor with MSE\n@@ -50,4 +51,4 @@\n         tensor(0.2500)\n     \"\"\"\n     sum_squared_error, n_obs = _mean_squared_error_update(preds, target)\n-    return _mean_squared_error_compute(sum_squared_error, n_obs)\n+    return _mean_squared_error_compute(sum_squared_error, n_obs, squared=squared)\ndiff --git a/torchmetrics/regression/mean_squared_error.py b/torchmetrics/regression/mean_squared_error.py\n--- a/torchmetrics/regression/mean_squared_error.py\n+++ b/torchmetrics/regression/mean_squared_error.py\n@@ -39,6 +39,8 @@\n             before returning the value at the step. default: False\n         process_group:\n             Specify the process group on which synchronization is called. default: None (which selects the entire world)\n+        squared:\n+            If True returns MSE value, if False returns RMSE value.\n \n     Example:\n         >>> from torchmetrics import MeanSquaredError\n@@ -56,6 +58,7 @@\n         dist_sync_on_step: bool = False,\n         process_group: Optional[Any] = None,\n         dist_sync_fn: Callable = None,\n+        squared: bool = True,\n     ):\n         super().__init__(\n             compute_on_step=compute_on_step,\n@@ -66,6 +69,7 @@\n \n         self.add_state(\"sum_squared_error\", default=tensor(0.0), dist_reduce_fx=\"sum\")\n         self.add_state(\"total\", default=tensor(0), dist_reduce_fx=\"sum\")\n+        self.squared = squared\n \n     def update(self, preds: Tensor, target: Tensor):\n         \"\"\"\n@@ -84,7 +88,7 @@\n         \"\"\"\n         Computes mean squared error over state.\n         \"\"\"\n-        return _mean_squared_error_compute(self.sum_squared_error, self.total)\n+        return _mean_squared_error_compute(self.sum_squared_error, self.total, squared=self.squared)\n \n     @property\n     def is_differentiable(self):\n", "issue": "Add option to compute root_mean_squared_error\n## \ud83d\ude80 Feature\r\nAllow the user to choose between MSE and RMSE.\r\n\r\n### Motivation\r\nIn a physical domain the RMSE, which is essentially the mean of distances, may be significantly more intuitive than the MSE. Therefore, it would be nice to have the option to choose the preferd metric.\r\n\r\n### Pitch\r\nSimilar to the implementation in [scikit-learn](https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_error.html#sklearn.metrics.mean_squared_error) one could simply pass `squared=False` to the `MeanSquaredError` module or the `mean_squared_error` function.\r\n\n", "before_files": [{"content": "# Copyright The PyTorch 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\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) -> Tuple[Tensor, int]:\n    _check_same_shape(preds, target)\n    diff = preds - target\n    sum_squared_error = torch.sum(diff * diff)\n    n_obs = target.numel()\n    return sum_squared_error, n_obs\n\n\ndef _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int) -> Tensor:\n    return sum_squared_error / n_obs\n\n\ndef mean_squared_error(preds: Tensor, target: Tensor) -> Tensor:\n    \"\"\"\n    Computes mean squared error\n\n    Args:\n        preds: estimated labels\n        target: ground truth labels\n\n    Return:\n        Tensor with MSE\n\n    Example:\n        >>> from torchmetrics.functional 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    sum_squared_error, n_obs = _mean_squared_error_update(preds, target)\n    return _mean_squared_error_compute(sum_squared_error, n_obs)\n", "path": "torchmetrics/functional/regression/mean_squared_error.py"}, {"content": "# Copyright The PyTorch 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 Any, Callable, Optional\n\nimport torch\nfrom torch import Tensor, tensor\n\nfrom torchmetrics.functional.regression.mean_squared_error import (\n    _mean_squared_error_compute,\n    _mean_squared_error_update,\n)\nfrom torchmetrics.metric import Metric\n\n\nclass MeanSquaredError(Metric):\n    r\"\"\"\n    Computes `mean squared error <https://en.wikipedia.org/wiki/Mean_squared_error>`_ (MSE):\n\n    .. math:: \\text{MSE} = \\frac{1}{N}\\sum_i^N(y_i - \\hat{y_i})^2\n\n    Where :math:`y` is a tensor of target values, and :math:`\\hat{y}` is a tensor of predictions.\n\n    Args:\n        compute_on_step:\n            Forward only calls ``update()`` and return None if this is set to False. default: True\n        dist_sync_on_step:\n            Synchronize metric state across processes at each ``forward()``\n            before returning the value at the step. default: False\n        process_group:\n            Specify the process group on which synchronization is called. default: None (which selects the entire world)\n\n    Example:\n        >>> from torchmetrics import MeanSquaredError\n        >>> target = torch.tensor([2.5, 5.0, 4.0, 8.0])\n        >>> preds = torch.tensor([3.0, 5.0, 2.5, 7.0])\n        >>> mean_squared_error = MeanSquaredError()\n        >>> mean_squared_error(preds, target)\n        tensor(0.8750)\n\n    \"\"\"\n\n    def __init__(\n        self,\n        compute_on_step: bool = True,\n        dist_sync_on_step: bool = False,\n        process_group: Optional[Any] = None,\n        dist_sync_fn: Callable = None,\n    ):\n        super().__init__(\n            compute_on_step=compute_on_step,\n            dist_sync_on_step=dist_sync_on_step,\n            process_group=process_group,\n            dist_sync_fn=dist_sync_fn,\n        )\n\n        self.add_state(\"sum_squared_error\", default=tensor(0.0), dist_reduce_fx=\"sum\")\n        self.add_state(\"total\", default=tensor(0), dist_reduce_fx=\"sum\")\n\n    def update(self, preds: Tensor, target: Tensor):\n        \"\"\"\n        Update state with predictions and targets.\n\n        Args:\n            preds: Predictions from model\n            target: Ground truth values\n        \"\"\"\n        sum_squared_error, n_obs = _mean_squared_error_update(preds, target)\n\n        self.sum_squared_error += sum_squared_error\n        self.total += n_obs\n\n    def compute(self):\n        \"\"\"\n        Computes mean squared error over state.\n        \"\"\"\n        return _mean_squared_error_compute(self.sum_squared_error, self.total)\n\n    @property\n    def is_differentiable(self):\n        return True\n", "path": "torchmetrics/regression/mean_squared_error.py"}], "after_files": [{"content": "# Copyright The PyTorch 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\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) -> Tuple[Tensor, int]:\n    _check_same_shape(preds, target)\n    diff = preds - target\n    sum_squared_error = torch.sum(diff * diff)\n    n_obs = target.numel()\n    return sum_squared_error, n_obs\n\n\ndef _mean_squared_error_compute(sum_squared_error: Tensor, n_obs: int, squared: bool = True) -> Tensor:\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) -> Tensor:\n    \"\"\"\n    Computes 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\n    Return:\n        Tensor with MSE\n\n    Example:\n        >>> from torchmetrics.functional 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    sum_squared_error, n_obs = _mean_squared_error_update(preds, target)\n    return _mean_squared_error_compute(sum_squared_error, n_obs, squared=squared)\n", "path": "torchmetrics/functional/regression/mean_squared_error.py"}, {"content": "# Copyright The PyTorch 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 Any, Callable, Optional\n\nimport torch\nfrom torch import Tensor, tensor\n\nfrom torchmetrics.functional.regression.mean_squared_error import (\n    _mean_squared_error_compute,\n    _mean_squared_error_update,\n)\nfrom torchmetrics.metric import Metric\n\n\nclass MeanSquaredError(Metric):\n    r\"\"\"\n    Computes `mean squared error <https://en.wikipedia.org/wiki/Mean_squared_error>`_ (MSE):\n\n    .. math:: \\text{MSE} = \\frac{1}{N}\\sum_i^N(y_i - \\hat{y_i})^2\n\n    Where :math:`y` is a tensor of target values, and :math:`\\hat{y}` is a tensor of predictions.\n\n    Args:\n        compute_on_step:\n            Forward only calls ``update()`` and return None if this is set to False. default: True\n        dist_sync_on_step:\n            Synchronize metric state across processes at each ``forward()``\n            before returning the value at the step. default: False\n        process_group:\n            Specify the process group on which synchronization is called. default: None (which selects the entire world)\n        squared:\n            If True returns MSE value, if False returns RMSE value.\n\n    Example:\n        >>> from torchmetrics import MeanSquaredError\n        >>> target = torch.tensor([2.5, 5.0, 4.0, 8.0])\n        >>> preds = torch.tensor([3.0, 5.0, 2.5, 7.0])\n        >>> mean_squared_error = MeanSquaredError()\n        >>> mean_squared_error(preds, target)\n        tensor(0.8750)\n\n    \"\"\"\n\n    def __init__(\n        self,\n        compute_on_step: bool = True,\n        dist_sync_on_step: bool = False,\n        process_group: Optional[Any] = None,\n        dist_sync_fn: Callable = None,\n        squared: bool = True,\n    ):\n        super().__init__(\n            compute_on_step=compute_on_step,\n            dist_sync_on_step=dist_sync_on_step,\n            process_group=process_group,\n            dist_sync_fn=dist_sync_fn,\n        )\n\n        self.add_state(\"sum_squared_error\", default=tensor(0.0), dist_reduce_fx=\"sum\")\n        self.add_state(\"total\", default=tensor(0), dist_reduce_fx=\"sum\")\n        self.squared = squared\n\n    def update(self, preds: Tensor, target: Tensor):\n        \"\"\"\n        Update state with predictions and targets.\n\n        Args:\n            preds: Predictions from model\n            target: Ground truth values\n        \"\"\"\n        sum_squared_error, n_obs = _mean_squared_error_update(preds, target)\n\n        self.sum_squared_error += sum_squared_error\n        self.total += n_obs\n\n    def compute(self):\n        \"\"\"\n        Computes mean squared error over state.\n        \"\"\"\n        return _mean_squared_error_compute(self.sum_squared_error, self.total, squared=self.squared)\n\n    @property\n    def is_differentiable(self):\n        return True\n", "path": "torchmetrics/regression/mean_squared_error.py"}]}