Repository: Lightning-AI/torchmetrics

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. -->



Discussion:
@SkafteNicki I can submit a PR if its something trivial.
Hi, I think we can replace `x.sqrt()` call with `x ** 0.5`, which can be also computed when passing `16bit` float tensor.
@shenoynikhil feel free to send a PR with the change proposed by @stancld :]
The fix gives this error,
`RuntimeError: "pow" not implemented for 'Half'`

So `sqrt()` does not have a problem with `bfloat16` but has a problem with `float16` ([source](https://github.com/pytorch/pytorch/issues/54774#issuecomment-808472386)). Do you think we can add something like to the `_pearson_corrcoef_compute` ,
```python
    # 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()
```

Resolving patch:
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 @@ def _pearson_corrcoef_compute(
     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)
 


Test patch:
diff --git a/tests/unittests/regression/test_pearson.py b/tests/unittests/regression/test_pearson.py
--- a/tests/unittests/regression/test_pearson.py
+++ b/tests/unittests/regression/test_pearson.py
@@ -100,8 +100,6 @@ def test_pearson_corrcoef_differentiability(self, preds, target):
             metric_functional=pearson_corrcoef,
         )
 
-    # Pearson half + cpu does not work due to missing support in torch.sqrt
-    @pytest.mark.xfail(reason="PearsonCorrCoef metric does not support cpu + half precision")
     def test_pearson_corrcoef_half_cpu(self, preds, target):
         """Test dtype support of the metric on CPU."""
         num_outputs = EXTRA_DIM if preds.ndim == 3 else 1