Repository: Lightning-AI/torchmetrics

Issue:
PearsonCorrCoeff returns nan when input is of type torch.float16 or torch.bfloat16.
## 🐛 Bug

PearsonCorrCoeff returns nan when input is of type torch.float16 or torch.bfloat16 and all values are close.

### To Reproduce

<!-- If you have a code sample, error messages, stack traces, please provide it here as well -->

<details>
  <summary>

```python
import torch
import torchmetrics as tm
pcc = tm.regression.PearsonCorrCoef().to("cuda")
pred = torch.tensor([0.4746, 0.4805, 0.4766, 0.4805, 0.4766, 0.4805, 0.4785, 0.4824, 0.4805],dtype=torch.float16).to("cuda")
target = torch.tensor([0.0336, 0.3676, 0.6302, 0.7192, 0.2295, 0.2886, 0.6302, 0.7096, 0.0208],dtype=torch.float16).to("cuda")
print(pcc(pred,target))
print(pcc(pred.to(torch.float32),target.to(torch.float32)))
tensor(nan, device='cuda:0')
tensor(0.3720, device='cuda:0')
```

</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

Python version: 3.10.9
Torch version: 1.12.1
TorchMetrics version: 0.11.1
GPU device name: Tesla T4
CUDA Version: 11.4 

### Additional context

When running in a training loop I found that some fraction (~30%) of steps would not produce a nan number when using torch.float16 or bfloat16, while the other ~70% would.
This seems to occur because the values in pred above are not very different (changing one value of pred above to be more different than the rest will compute a correct PCC), however I think that this should still be able to be computed with half precision and the standard deviation of pred shown above.
<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!

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
@@ -49,13 +49,25 @@ def _pearson_corrcoef_update(
     # Data checking
     _check_same_shape(preds, target)
     _check_data_shape_to_num_outputs(preds, target, num_outputs)
+    cond = n_prior.mean() > 0
 
     n_obs = preds.shape[0]
-    mx_new = (n_prior * mean_x + preds.mean(0) * n_obs) / (n_prior + n_obs)
-    my_new = (n_prior * mean_y + target.mean(0) * n_obs) / (n_prior + n_obs)
+    if cond:
+        mx_new = (n_prior * mean_x + preds.sum(0)) / (n_prior + n_obs)
+        my_new = (n_prior * mean_y + target.sum(0)) / (n_prior + n_obs)
+    else:
+        mx_new = preds.mean(0)
+        my_new = target.mean(0)
+
     n_prior += n_obs
-    var_x += ((preds - mx_new) * (preds - mean_x)).sum(0)
-    var_y += ((target - my_new) * (target - mean_y)).sum(0)
+
+    if cond:
+        var_x += ((preds - mx_new) * (preds - mean_x)).sum(0)
+        var_y += ((target - my_new) * (target - mean_y)).sum(0)
+
+    else:
+        var_x += preds.var(0) * (n_obs - 1)
+        var_y += target.var(0) * (n_obs - 1)
     corr_xy += ((preds - mx_new) * (target - mean_y)).sum(0)
     mean_x = mx_new
     mean_y = my_new


Test patch: