Allow sharding for distributed tests
Addresses my mistake introduced in https://github.com/pytorch/pytorch/pull/76536#issuecomment-1112657429
Also allows for sharding 1 in run_test.py
Pull Request resolved: https://github.com/pytorch/pytorch/pull/76570
Approved by: https://github.com/jeffdaily, https://github.com/seemethere
