We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
PCC gives error with AMP
## 🐛 Bug

On turning on Automatic Mixed Precision, there is an error in Pearson Correlation Coefficient computation at this [line](https://github.com/Lightning-AI/torchmetrics/blob/962f82db7b0ab7ca32fe050aba45f1d3f9fe64c9/src/torchmetrics/functional/regression/pearson.py#L95). The trace is as follows,
```
  File "/home/nikhil/miniconda3/envs/mrl/lib/python3.8/site-packages/torchmetrics/functional/regression/pearson.py", line 80, in _pearson_corrcoef_compute                                              
    corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()                                                                                                                                             
RuntimeError: "sqrt_vml_cpu" not implemented for 'Half'
```

Turning AMP off, it fixes it. 

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

Steps to reproduce the behavior...

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details>
  <summary>Code sample</summary>

<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

</details>

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

- TorchMetrics version : `0.11.4` installed via `pip`
- Python & PyTorch Version (e.g., 1.0): `Python 3.8`, `Torch 2.0.1+cu118`
- Any other relevant information such as OS (e.g., Linux): `Linux`

### Additional context

<!-- Add any other context about the problem here. -->


--- 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/pearson.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
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 _pearson_corrcoef_update(
24     preds: Tensor,
25     target: Tensor,
26     mean_x: Tensor,
27     mean_y: Tensor,
28     var_x: Tensor,
29     var_y: Tensor,
30     corr_xy: Tensor,
31     n_prior: Tensor,
32     num_outputs: int,
33 ) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
34     """Update and returns variables required to compute Pearson Correlation Coefficient.
35 
36     Check for same shape of input tensors.
37 
38     Args:
39         preds: estimated scores
40         target: ground truth scores
41         mean_x: current mean estimate of x tensor
42         mean_y: current mean estimate of y tensor
43         var_x: current variance estimate of x tensor
44         var_y: current variance estimate of y tensor
45         corr_xy: current covariance estimate between x and y tensor
46         n_prior: current number of observed observations
47         num_outputs: Number of outputs in multioutput setting
48     """
49     # Data checking
50     _check_same_shape(preds, target)
51     _check_data_shape_to_num_outputs(preds, target, num_outputs)
52     cond = n_prior.mean() > 0
53 
54     n_obs = preds.shape[0]
55     if cond:
56         mx_new = (n_prior * mean_x + preds.sum(0)) / (n_prior + n_obs)
57         my_new = (n_prior * mean_y + target.sum(0)) / (n_prior + n_obs)
58     else:
59         mx_new = preds.mean(0)
60         my_new = target.mean(0)
61 
62     n_prior += n_obs
63 
64     if cond:
65         var_x += ((preds - mx_new) * (preds - mean_x)).sum(0)
66         var_y += ((target - my_new) * (target - mean_y)).sum(0)
67 
68     else:
69         var_x += preds.var(0) * (n_obs - 1)
70         var_y += target.var(0) * (n_obs - 1)
71     corr_xy += ((preds - mx_new) * (target - mean_y)).sum(0)
72     mean_x = mx_new
73     mean_y = my_new
74 
75     return mean_x, mean_y, var_x, var_y, corr_xy, n_prior
76 
77 
78 def _pearson_corrcoef_compute(
79     var_x: Tensor,
80     var_y: Tensor,
81     corr_xy: Tensor,
82     nb: Tensor,
83 ) -> Tensor:
84     """Compute the final pearson correlation based on accumulated statistics.
85 
86     Args:
87         var_x: variance estimate of x tensor
88         var_y: variance estimate of y tensor
89         corr_xy: covariance estimate between x and y tensor
90         nb: number of observations
91     """
92     var_x /= nb - 1
93     var_y /= nb - 1
94     corr_xy /= nb - 1
95     corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()
96     return torch.clamp(corrcoef, -1.0, 1.0)
97 
98 
99 def pearson_corrcoef(preds: Tensor, target: Tensor) -> Tensor:
100     """Compute pearson correlation coefficient.
101 
102     Args:
103         preds: estimated scores
104         target: ground truth scores
105 
106     Example (single output regression):
107         >>> from torchmetrics.functional.regression import pearson_corrcoef
108         >>> target = torch.tensor([3, -0.5, 2, 7])
109         >>> preds = torch.tensor([2.5, 0.0, 2, 8])
110         >>> pearson_corrcoef(preds, target)
111         tensor(0.9849)
112 
113     Example (multi output regression):
114         >>> from torchmetrics.functional.regression import pearson_corrcoef
115         >>> target = torch.tensor([[3, -0.5], [2, 7]])
116         >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])
117         >>> pearson_corrcoef(preds, target)
118         tensor([1., 1.])
119     """
120     d = preds.shape[1] if preds.ndim == 2 else 1
121     _temp = torch.zeros(d, dtype=preds.dtype, device=preds.device)
122     mean_x, mean_y, var_x = _temp.clone(), _temp.clone(), _temp.clone()
123     var_y, corr_xy, nb = _temp.clone(), _temp.clone(), _temp.clone()
124     _, _, var_x, var_y, corr_xy, nb = _pearson_corrcoef_update(
125         preds, target, mean_x, mean_y, var_x, var_y, corr_xy, nb, num_outputs=1 if preds.ndim == 1 else preds.shape[-1]
126     )
127     return _pearson_corrcoef_compute(var_x, var_y, corr_xy, nb)
128
```
--- 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/pearson.py b/src/torchmetrics/functional/regression/pearson.py
--- a/src/torchmetrics/functional/regression/pearson.py
+++ b/src/torchmetrics/functional/regression/pearson.py
@@ -92,6 +92,12 @@
     var_x /= nb - 1
     var_y /= nb - 1
     corr_xy /= nb - 1
+    # if var_x, var_y is float16 and on cpu, make it bfloat16 as sqrt is not supported for float16
+    # on cpu, remove this after https://github.com/pytorch/pytorch/issues/54774 is fixed
+    if var_x.dtype == torch.float16 and var_x.device == torch.device("cpu"):
+        var_x = var_x.bfloat16()
+        var_y = var_y.bfloat16()
+
     corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()
     return torch.clamp(corrcoef, -1.0, 1.0)

{"golden_diff": "diff --git a/src/torchmetrics/functional/regression/pearson.py b/src/torchmetrics/functional/regression/pearson.py\n--- a/src/torchmetrics/functional/regression/pearson.py\n+++ b/src/torchmetrics/functional/regression/pearson.py\n@@ -92,6 +92,12 @@\n     var_x /= nb - 1\n     var_y /= nb - 1\n     corr_xy /= nb - 1\n+    # if var_x, var_y is float16 and on cpu, make it bfloat16 as sqrt is not supported for float16\n+    # on cpu, remove this after https://github.com/pytorch/pytorch/issues/54774 is fixed\n+    if var_x.dtype == torch.float16 and var_x.device == torch.device(\"cpu\"):\n+        var_x = var_x.bfloat16()\n+        var_y = var_y.bfloat16()\n+\n     corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()\n     return torch.clamp(corrcoef, -1.0, 1.0)\n", "issue": "PCC gives error with AMP\n## \ud83d\udc1b Bug\r\n\r\nOn turning on Automatic Mixed Precision, there is an error in Pearson Correlation Coefficient computation at this [line](https://github.com/Lightning-AI/torchmetrics/blob/962f82db7b0ab7ca32fe050aba45f1d3f9fe64c9/src/torchmetrics/functional/regression/pearson.py#L95). The trace is as follows,\r\n```\r\n  File \"/home/nikhil/miniconda3/envs/mrl/lib/python3.8/site-packages/torchmetrics/functional/regression/pearson.py\", line 80, in _pearson_corrcoef_compute                                              \r\n    corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()                                                                                                                                             \r\nRuntimeError: \"sqrt_vml_cpu\" not implemented for 'Half'\r\n```\r\n\r\nTurning AMP off, it fixes it. \r\n\r\n<!-- A clear and concise description of what the bug is. -->\r\n\r\n### To Reproduce\r\n\r\nSteps to reproduce the behavior...\r\n\r\n<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->\r\n\r\n<details>\r\n  <summary>Code sample</summary>\r\n\r\n<!-- Ideally attach a minimal code sample to reproduce the decried issue.\r\nMinimal means having the shortest code but still preserving the bug. -->\r\n\r\n</details>\r\n\r\n### Expected behavior\r\n\r\n<!-- A clear and concise description of what you expected to happen. -->\r\n\r\n### Environment\r\n\r\n- TorchMetrics version : `0.11.4` installed via `pip`\r\n- Python & PyTorch Version (e.g., 1.0): `Python 3.8`, `Torch 2.0.1+cu118`\r\n- Any other relevant information such as OS (e.g., Linux): `Linux`\r\n\r\n### Additional context\r\n\r\n<!-- Add any other context about the problem here. -->\r\n\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\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 _pearson_corrcoef_update(\n    preds: Tensor,\n    target: Tensor,\n    mean_x: Tensor,\n    mean_y: Tensor,\n    var_x: Tensor,\n    var_y: Tensor,\n    corr_xy: Tensor,\n    n_prior: Tensor,\n    num_outputs: int,\n) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:\n    \"\"\"Update and returns variables required to compute Pearson Correlation Coefficient.\n\n    Check for same shape of input tensors.\n\n    Args:\n        preds: estimated scores\n        target: ground truth scores\n        mean_x: current mean estimate of x tensor\n        mean_y: current mean estimate of y tensor\n        var_x: current variance estimate of x tensor\n        var_y: current variance estimate of y tensor\n        corr_xy: current covariance estimate between x and y tensor\n        n_prior: current number of observed observations\n        num_outputs: Number of outputs in multioutput setting\n    \"\"\"\n    # Data checking\n    _check_same_shape(preds, target)\n    _check_data_shape_to_num_outputs(preds, target, num_outputs)\n    cond = n_prior.mean() > 0\n\n    n_obs = preds.shape[0]\n    if cond:\n        mx_new = (n_prior * mean_x + preds.sum(0)) / (n_prior + n_obs)\n        my_new = (n_prior * mean_y + target.sum(0)) / (n_prior + n_obs)\n    else:\n        mx_new = preds.mean(0)\n        my_new = target.mean(0)\n\n    n_prior += n_obs\n\n    if cond:\n        var_x += ((preds - mx_new) * (preds - mean_x)).sum(0)\n        var_y += ((target - my_new) * (target - mean_y)).sum(0)\n\n    else:\n        var_x += preds.var(0) * (n_obs - 1)\n        var_y += target.var(0) * (n_obs - 1)\n    corr_xy += ((preds - mx_new) * (target - mean_y)).sum(0)\n    mean_x = mx_new\n    mean_y = my_new\n\n    return mean_x, mean_y, var_x, var_y, corr_xy, n_prior\n\n\ndef _pearson_corrcoef_compute(\n    var_x: Tensor,\n    var_y: Tensor,\n    corr_xy: Tensor,\n    nb: Tensor,\n) -> Tensor:\n    \"\"\"Compute the final pearson correlation based on accumulated statistics.\n\n    Args:\n        var_x: variance estimate of x tensor\n        var_y: variance estimate of y tensor\n        corr_xy: covariance estimate between x and y tensor\n        nb: number of observations\n    \"\"\"\n    var_x /= nb - 1\n    var_y /= nb - 1\n    corr_xy /= nb - 1\n    corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()\n    return torch.clamp(corrcoef, -1.0, 1.0)\n\n\ndef pearson_corrcoef(preds: Tensor, target: Tensor) -> Tensor:\n    \"\"\"Compute pearson correlation coefficient.\n\n    Args:\n        preds: estimated scores\n        target: ground truth scores\n\n    Example (single output regression):\n        >>> from torchmetrics.functional.regression import pearson_corrcoef\n        >>> target = torch.tensor([3, -0.5, 2, 7])\n        >>> preds = torch.tensor([2.5, 0.0, 2, 8])\n        >>> pearson_corrcoef(preds, target)\n        tensor(0.9849)\n\n    Example (multi output regression):\n        >>> from torchmetrics.functional.regression import pearson_corrcoef\n        >>> target = torch.tensor([[3, -0.5], [2, 7]])\n        >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])\n        >>> pearson_corrcoef(preds, target)\n        tensor([1., 1.])\n    \"\"\"\n    d = preds.shape[1] if preds.ndim == 2 else 1\n    _temp = torch.zeros(d, dtype=preds.dtype, device=preds.device)\n    mean_x, mean_y, var_x = _temp.clone(), _temp.clone(), _temp.clone()\n    var_y, corr_xy, nb = _temp.clone(), _temp.clone(), _temp.clone()\n    _, _, var_x, var_y, corr_xy, nb = _pearson_corrcoef_update(\n        preds, target, mean_x, mean_y, var_x, var_y, corr_xy, nb, num_outputs=1 if preds.ndim == 1 else preds.shape[-1]\n    )\n    return _pearson_corrcoef_compute(var_x, var_y, corr_xy, nb)\n", "path": "src/torchmetrics/functional/regression/pearson.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\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 _pearson_corrcoef_update(\n    preds: Tensor,\n    target: Tensor,\n    mean_x: Tensor,\n    mean_y: Tensor,\n    var_x: Tensor,\n    var_y: Tensor,\n    corr_xy: Tensor,\n    n_prior: Tensor,\n    num_outputs: int,\n) -> Tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:\n    \"\"\"Update and returns variables required to compute Pearson Correlation Coefficient.\n\n    Check for same shape of input tensors.\n\n    Args:\n        preds: estimated scores\n        target: ground truth scores\n        mean_x: current mean estimate of x tensor\n        mean_y: current mean estimate of y tensor\n        var_x: current variance estimate of x tensor\n        var_y: current variance estimate of y tensor\n        corr_xy: current covariance estimate between x and y tensor\n        n_prior: current number of observed observations\n        num_outputs: Number of outputs in multioutput setting\n    \"\"\"\n    # Data checking\n    _check_same_shape(preds, target)\n    _check_data_shape_to_num_outputs(preds, target, num_outputs)\n    cond = n_prior.mean() > 0\n\n    n_obs = preds.shape[0]\n    if cond:\n        mx_new = (n_prior * mean_x + preds.sum(0)) / (n_prior + n_obs)\n        my_new = (n_prior * mean_y + target.sum(0)) / (n_prior + n_obs)\n    else:\n        mx_new = preds.mean(0)\n        my_new = target.mean(0)\n\n    n_prior += n_obs\n\n    if cond:\n        var_x += ((preds - mx_new) * (preds - mean_x)).sum(0)\n        var_y += ((target - my_new) * (target - mean_y)).sum(0)\n\n    else:\n        var_x += preds.var(0) * (n_obs - 1)\n        var_y += target.var(0) * (n_obs - 1)\n    corr_xy += ((preds - mx_new) * (target - mean_y)).sum(0)\n    mean_x = mx_new\n    mean_y = my_new\n\n    return mean_x, mean_y, var_x, var_y, corr_xy, n_prior\n\n\ndef _pearson_corrcoef_compute(\n    var_x: Tensor,\n    var_y: Tensor,\n    corr_xy: Tensor,\n    nb: Tensor,\n) -> Tensor:\n    \"\"\"Compute the final pearson correlation based on accumulated statistics.\n\n    Args:\n        var_x: variance estimate of x tensor\n        var_y: variance estimate of y tensor\n        corr_xy: covariance estimate between x and y tensor\n        nb: number of observations\n    \"\"\"\n    var_x /= nb - 1\n    var_y /= nb - 1\n    corr_xy /= nb - 1\n    # if var_x, var_y is float16 and on cpu, make it bfloat16 as sqrt is not supported for float16\n    # on cpu, remove this after https://github.com/pytorch/pytorch/issues/54774 is fixed\n    if var_x.dtype == torch.float16 and var_x.device == torch.device(\"cpu\"):\n        var_x = var_x.bfloat16()\n        var_y = var_y.bfloat16()\n\n    corrcoef = (corr_xy / (var_x * var_y).sqrt()).squeeze()\n    return torch.clamp(corrcoef, -1.0, 1.0)\n\n\ndef pearson_corrcoef(preds: Tensor, target: Tensor) -> Tensor:\n    \"\"\"Compute pearson correlation coefficient.\n\n    Args:\n        preds: estimated scores\n        target: ground truth scores\n\n    Example (single output regression):\n        >>> from torchmetrics.functional.regression import pearson_corrcoef\n        >>> target = torch.tensor([3, -0.5, 2, 7])\n        >>> preds = torch.tensor([2.5, 0.0, 2, 8])\n        >>> pearson_corrcoef(preds, target)\n        tensor(0.9849)\n\n    Example (multi output regression):\n        >>> from torchmetrics.functional.regression import pearson_corrcoef\n        >>> target = torch.tensor([[3, -0.5], [2, 7]])\n        >>> preds = torch.tensor([[2.5, 0.0], [2, 8]])\n        >>> pearson_corrcoef(preds, target)\n        tensor([1., 1.])\n    \"\"\"\n    d = preds.shape[1] if preds.ndim == 2 else 1\n    _temp = torch.zeros(d, dtype=preds.dtype, device=preds.device)\n    mean_x, mean_y, var_x = _temp.clone(), _temp.clone(), _temp.clone()\n    var_y, corr_xy, nb = _temp.clone(), _temp.clone(), _temp.clone()\n    _, _, var_x, var_y, corr_xy, nb = _pearson_corrcoef_update(\n        preds, target, mean_x, mean_y, var_x, var_y, corr_xy, nb, num_outputs=1 if preds.ndim == 1 else preds.shape[-1]\n    )\n    return _pearson_corrcoef_compute(var_x, var_y, corr_xy, nb)\n", "path": "src/torchmetrics/functional/regression/pearson.py"}]}