Skip to content

fix(torch-stubs): declare torch.device as a context manager - #4912

Closed
yeetypete wants to merge 1 commit into
facebook:mainfrom
yeetypete:torch-stubs-device-context-manager
Closed

yeetypete wants to merge 1 commit into
facebook:mainfrom
yeetypete:torch-stubs-device-context-manager

Conversation

@yeetypete

Copy link
Copy Markdown
Contributor

Summary

Since PyTorch 2.0, torch.device can be used as a context manager to set the default device for a block. The shape stubs' device had no __enter__/__exit__, so every with torch.device(...) block reported bad-context-manager. We add them here, matching the PyTorch stubs:

https://github.com/pytorch/pytorch/blob/d6c03540dc0f040e1d8a920d19d9333de6513545/torch/_C/__init__.pyi.in#L121-L122

Test Plan

Run test.py.

@meta-codesync

meta-codesync Bot commented Sep 12, 2026

Copy link
Copy Markdown
Contributor

This pull request has been imported. If you are a Meta employee, you can view this in D119797451. (Because this pull request was imported automatically, there will not be any future comments.)

@stroxler stroxler left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Review automatically exported from Phabricator review in Meta.

@meta-codesync meta-codesync Bot closed this in 9f07622 Sep 14, 2026
@meta-codesync meta-codesync Bot added the Merged label Sep 14, 2026
@meta-codesync

meta-codesync Bot commented Sep 14, 2026

Copy link
Copy Markdown
Contributor

This pull request has been merged in 9f07622.

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

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants