Skip to content

[wip] fix(metric): honor torch.device context manager when setting default device - #3448

Open
bhimrazy wants to merge 2 commits into
Lightning-AI:masterfrom
bhimrazy:fix/metric-default-device-torch-2.8
Open

[wip] fix(metric): honor torch.device context manager when setting default device#3448
bhimrazy wants to merge 2 commits into
Lightning-AI:masterfrom
bhimrazy:fix/metric-default-device-torch-2.8

Conversation

@bhimrazy

@bhimrazy bhimrazy commented Aug 9, 2026

Copy link
Copy Markdown
Collaborator

What does this PR do?

Metric.__init__ gates torch.get_default_device() on torch >= 2.3, but that function only started reflecting an active with torch.device(...) context manager in 2.8 — before that it tracked only set_default_device(). So on torch 2.3–2.7, a metric built inside such a block gets CPU instead of the context device.

Gates on 2.8 and keeps the torch.empty(0).device probe below it, which does respect the context manager.

Verified on torch 2.6.0 inside with torch.device("meta"): get_default_device()cpu, torch.empty(0).devicemeta.

The existing coverage for this path is GPU-only and passes today only because the GPU images are torch 2.0 and 2.8; any image in 2.3–2.7 turns it red. Adds a CPU-runnable regression test using the meta device.

Prerequisite for #3449 (floor moves to 2.6).

…t device

`Metric.__init__` gated the use of `torch.get_default_device()` on torch >= 2.3.
However, until torch 2.8 that function only reflected `set_default_device()` and
ignored an active `with torch.device(...)` context manager, so metrics created
inside such a block were assigned the CPU instead of the context device.

torch 2.8 (pytorch/pytorch v2.8.0) added the `DeviceContext` lookup to
`get_default_device`, so gate on 2.8 and keep the `torch.empty(0).device` probe
for older versions, which does respect the context manager.

Verified on torch 2.6.0: inside `with torch.device("meta")`,
`get_default_device()` returns cpu while `torch.empty(0).device` returns meta.

Adds a CPU-runnable regression test using the meta device; the existing coverage
for this path was GPU-only.
@codecov-commenter

codecov-commenter commented Aug 9, 2026

Copy link
Copy Markdown

⚠️ Please install the 'codecov app svg image' to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 36%. Comparing base (d184220) to head (166e3a6).
⚠️ Report is 10 commits behind head on master.
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files
@@          Coverage Diff           @@
##           master   #3448   +/-   ##
======================================
- Coverage      37%     36%   -0%     
======================================
  Files         349     349           
  Lines       19901   19910    +9     
======================================
+ Hits         7264    7265    +1     
- Misses      12637   12645    +8     
🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants