Better error message on wrong device

CHANGELOG.md

@@ -11,7 +11,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
--
+- Added specific runtime error when metric object is on wrong device (#1056(https://github.com/PyTorchLightning/metrics/pull/1056))
 
 -
 

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

@@ -377,7 +377,18 @@ 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(
+ 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."
+ ) from err
+ raise err
+
  if self.compute_on_cpu:
  self._move_list_states_to_cpu()