fix bug for multigpus when only a subset is visible (#1332)
Summary:
When there are multiple GPUs and only a subset is visible, e.g. specified the visible GPUs by `CUDA_VISIBLE_DEVICES`, to the process, errors occur as described in https://github.com/pytorch/benchmark/issues/1331
fixes https://github.com/pytorch/benchmark/issues/1331
Pull Request resolved: https://github.com/pytorch/benchmark/pull/1332
Reviewed By: xuzhao9
Differential Revision: D41628139
Pulled By: FindHao
fbshipit-source-id: 284e9467d187c6ba194d48a106cb2140d091836b
