Repository: Lightning-AI/torchmetrics

Issue:
PearsonCorrCoef throws device error when tensors are on cuda device
## 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

### To Reproduce

```pthon
m = PearsonCorrCoef()
t1 = torch.tensor(range(10), device="cuda", dtype=float)
t2 = 2 * torch.tensor(range(10), device="cuda", dtype=float)
m(t1, t2)
```
### Expected behavior

Compite pesron corr coef

- TorchMetrics version: installed with pip , verison 0.8.2


Discussion:
Hi! thanks for your contribution!, great first issue!
As described [here in the documentation](https://torchmetrics.readthedocs.io/en/stable/pages/overview.html#metrics-and-devices) each metric object is a `nn.Module` and therefore should be moved to the same device as input. Your example therefore works if we move the metric to gpu:
```python
m = PearsonCorrCoef().cuda()
t1 = torch.tensor(range(10), device="cuda", dtype=float)
t2 = 2 * torch.tensor(range(10), device="cuda", dtype=float)
m(t1, t2)
```

Resolving patch:
diff --git a/torchmetrics/metric.py b/torchmetrics/metric.py
--- a/torchmetrics/metric.py
+++ b/torchmetrics/metric.py
@@ -378,7 +378,20 @@ def wrapped_func(*args: Any, **kwargs: Any) -> None:
             self._computed = None
             self._update_count += 1
             with torch.set_grad_enabled(self._enable_grad):
-                update(*args, **kwargs)
+                try:
+                    update(*args, **kwargs)
+                except RuntimeError as err:
+                    if "Expected all tensors to be on" in str(err):
+                        raise RuntimeError(
+                            "Encountered different devices in metric calculation"
+                            " (see stacktrace for details)."
+                            "This could be due to the metric class not being on the same device as input."
+                            f"Instead of `metric={self.__class__.__name__}(...)` try to do"
+                            f" `metric={self.__class__.__name__}(...).to(device)` where"
+                            " device corresponds to the device of the input."
+                        ) from err
+                    raise err
+
             if self.compute_on_cpu:
                 self._move_list_states_to_cpu()
 


Test patch:
diff --git a/tests/bases/test_metric.py b/tests/bases/test_metric.py
--- a/tests/bases/test_metric.py
+++ b/tests/bases/test_metric.py
@@ -26,6 +26,7 @@
 from tests.helpers import seed_all
 from tests.helpers.testers import DummyListMetric, DummyMetric, DummyMetricMultiOutput, DummyMetricSum
 from tests.helpers.utilities import no_warning_call
+from torchmetrics import PearsonCorrCoef
 from torchmetrics.utilities.imports import _TORCH_LOWER_1_6
 
 seed_all(42)
@@ -426,6 +427,17 @@ class UnsetProperty(metric_class):
         UnsetProperty()
 
 
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires gpu")
+def test_specific_error_on_wrong_device():
+    metric = PearsonCorrCoef()
+    preds = torch.tensor(range(10), device="cuda", dtype=torch.float)
+    target = torch.tensor(range(10), device="cuda", dtype=torch.float)
+    with pytest.raises(
+        RuntimeError, match="This could be due to the metric class not being on the same device as input"
+    ):
+        _ = metric(preds, target)
+
+
 @pytest.mark.parametrize("metric_class", [DummyListMetric, DummyMetric, DummyMetricMultiOutput, DummyMetricSum])
 def test_no_warning_on_custom_forward(metric_class):
     """If metric is using custom forward, full_state_update is irrelevant."""