Change current device in stream context manager if necessary#16128
Change current device in stream context manager if necessary#16128mrshenli wants to merge 3 commits into
Conversation
… stream context manager
facebook-github-bot
left a comment
There was a problem hiding this comment.
@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
facebook-github-bot
left a comment
There was a problem hiding this comment.
@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
|
|
||
| @unittest.skipIf(not TEST_MULTIGPU, "detected only one GPU") | ||
| @skipIfRocm | ||
| def test_stream_context(self): |
There was a problem hiding this comment.
I would prefer a more targeted test-case here for the stream context behavior. For example, this implicitly tests that the device and stream switch by using query() and _sleep. I'd rather just check that the device and stream are set appropriately and reset appropriately. For example, something like:
- Switch to s1, check that torch.cuda.current_stream() is s1 and current_device() is 1
- Check that stream and device are restored after context manager
- Check nested context managers
facebook-github-bot
left a comment
There was a problem hiding this comment.
@mrshenli has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.
…#16128) Summary: Fixes pytorch#16019 Pull Request resolved: pytorch#16128 Differential Revision: D13721850 Pulled By: mrshenli fbshipit-source-id: 422c6c0b97c1cd46e127e265b532cb8c74a3aac5
Fixes #16019