[sharded_tensor] fix typing issue for placement (#63426)
authorWanchao Liang <wanchaol@fb.com>
Wed, 18 Aug 2021 06:10:48 +0000 (23:10 -0700)
committerFacebook GitHub Bot <facebook-github-bot@users.noreply.github.com>
Wed, 18 Aug 2021 06:11:48 +0000 (23:11 -0700)
commitd431c77d76b8c92dfddfdfec2ce29b2e46c441c2
tree5bf97407bd58e73407b0bd3eec20084f394c2481
parent2fd14735d677d1c2cf87e4c76aea2111bc30c17b
[sharded_tensor] fix typing issue for placement (#63426)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/63426

placement should either be a string or a _remote_device, this fixes the type to match the behaviors
ghstack-source-id: 136041125

Reviewed By: pritamdamania87

Differential Revision: D30379702

fbshipit-source-id: 34e226494240923b433e3a39cc08c84d42cdad6b
torch/distributed/_sharding_spec/_internals.py