diff --git a/CHANGELOG.md b/CHANGELOG.md index 9518750113a..d33d477ef8d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -36,6 +36,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Fixed malformed LaTeX in `CLIPScore` and `HausdorffDistance` docstring math so it renders correctly ([#3427](https://github.com/Lightning-AI/torchmetrics/pull/3427)) +- Fixed `Metric` ignoring an active `torch.device` context manager on torch 2.3-2.7 ([#3448](https://github.com/Lightning-AI/torchmetrics/pull/3448)) + + --- ## [1.9.0] - 2026-03-05 diff --git a/src/torchmetrics/metric.py b/src/torchmetrics/metric.py index 238c3e5d112..6464a227019 100644 --- a/src/torchmetrics/metric.py +++ b/src/torchmetrics/metric.py @@ -39,7 +39,7 @@ ) from torchmetrics.utilities.distributed import gather_all_tensors from torchmetrics.utilities.exceptions import TorchMetricsUserError -from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_1, _TORCH_GREATER_EQUAL_2_3 +from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_1, _TORCH_GREATER_EQUAL_2_8 from torchmetrics.utilities.plot import _AX_TYPE, _PLOT_OUT_TYPE, plot_single_or_multi_val from torchmetrics.utilities.prints import rank_zero_warn @@ -113,7 +113,8 @@ def __init__( torch._C._log_api_usage_once(f"torchmetrics.metric.{self.__class__.__name__}") # magic patch for `RuntimeError: DataLoader worker (pid(s) 104) exited unexpectedly` self._TORCH_GREATER_EQUAL_2_1 = bool(_TORCH_GREATER_EQUAL_2_1) - self._device = torch.get_default_device() if _TORCH_GREATER_EQUAL_2_3 else torch.empty(0).device + # before 2.8, `get_default_device` only tracked `set_default_device`, not an active `torch.device` context + self._device = torch.get_default_device() if _TORCH_GREATER_EQUAL_2_8 else torch.empty(0).device self._dtype = torch.get_default_dtype() self.compute_on_cpu = kwargs.pop("compute_on_cpu", False) diff --git a/src/torchmetrics/utilities/imports.py b/src/torchmetrics/utilities/imports.py index 88da9076269..c5dd3796c5c 100644 --- a/src/torchmetrics/utilities/imports.py +++ b/src/torchmetrics/utilities/imports.py @@ -23,6 +23,7 @@ _TORCH_GREATER_EQUAL_2_2 = RequirementCache("torch>=2.2.0") _TORCH_GREATER_EQUAL_2_3 = RequirementCache("torch>=2.3.0") _TORCH_GREATER_EQUAL_2_5 = RequirementCache("torch>=2.5.0") +_TORCH_GREATER_EQUAL_2_8 = RequirementCache("torch>=2.8.0") _TORCH_LESS_THAN_2_6 = RequirementCache("torch<2.6.0") _TORCHMETRICS_GREATER_EQUAL_1_6 = RequirementCache("torchmetrics>=1.7.0") diff --git a/tests/unittests/bases/test_metric.py b/tests/unittests/bases/test_metric.py index 888dc6d864d..e4affe37449 100644 --- a/tests/unittests/bases/test_metric.py +++ b/tests/unittests/bases/test_metric.py @@ -30,6 +30,7 @@ from torchmetrics.clustering import AdjustedRandScore from torchmetrics.image import StructuralSimilarityIndexMeasure from torchmetrics.regression import PearsonCorrCoef, R2Score +from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_2_8 from unittests._helpers import seed_all from unittests._helpers.testers import DummyListMetric, DummyMetric, DummyMetricMultiOutput, DummyMetricSum @@ -333,6 +334,22 @@ def test_device_and_dtype_transfer(tmpdir): torch.set_default_dtype(default_dtype) +@pytest.mark.parametrize( + "use_get_default_device", + [ + pytest.param(True, marks=pytest.mark.skipif(not _TORCH_GREATER_EQUAL_2_8, reason="requires torch>=2.8")), + False, + ], +) +def test_device_from_device_context_manager(monkeypatch, use_get_default_device): + """Test that a metric picks up the device from an active `torch.device` context manager, on both version paths.""" + monkeypatch.setattr("torchmetrics.metric._TORCH_GREATER_EQUAL_2_8", use_get_default_device) + with torch.device("meta"): + metric = DummyMetricSum() + assert metric.device == torch.device("meta") + assert metric.x.device == torch.device("meta") + + def test_disable_of_normal_dtype_methods(): """Check that the default dtype changing methods does nothing.""" metric = DummyMetricSum()