Change current device in stream context manager if necessary (#16128)
Summary:
Fixes #16019
Pull Request resolved: https://github.com/pytorch/pytorch/pull/16128
Differential Revision: D13721850
Pulled By: mrshenli
fbshipit-source-id: 422c6c0b97c1cd46e127e265b532cb8c74a3aac5
