Better error message on wrong device

tests/bases/test_metric.py

@@ -25,6 +25,7 @@
 
 from tests.helpers import seed_all
 from tests.helpers.testers import DummyListMetric, DummyMetric, DummyMetricMultiOutput, DummyMetricSum
+from torchmetrics import PearsonCorrCoef
 from torchmetrics.utilities.imports import _TORCH_LOWER_1_6
 
 seed_all(42)
@@ -423,3 +424,14 @@ class UnsetProperty(metric_class):
         match="Torchmetrics v0.9 introduced a new argument class property called.*",
     ):
         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)

torchmetrics/metric.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 import functools
 import inspect
+import traceback
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
 from copy import deepcopy
@@ -377,7 +378,19 @@ 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):
+                        traceback.print_exc()
+                        raise RuntimeError(
+                            f"{str(err)}. \n"
+                            "This could be due to the metric class not being on the same device as input.\n"
+                            "Instead of `metric=Metric()` try to do `metric=Metric().to(device)` where"
+                            " device corresponds to the device of the input."
+                        )
+                    raise err
+
             if self.compute_on_cpu:
                 self._move_list_states_to_cpu()