Add better device idx parse checks (#37376)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/32079
Pull Request resolved: https://github.com/pytorch/pytorch/pull/37376
Differential Revision: D21476036
Pulled By: zou3519
fbshipit-source-id: 86907083c23cbaf165b645307fb340f2656b814e
