Fix BN tests for >= 8 GPU test environments (pytorch#19049)

Summary:
DDP does not support replicating BN layers within a process. Existing BN tests fail if the test environment has more than 8 GPUs. This is fixed by explicitly setting each process to use a single replica.
Pull Request resolved: pytorch#19049

Differential Revision: D14845286

Pulled By: mrshenli

fbshipit-source-id: 937dda5081d415ece48b21f2781b6b4e008dd42f

test/test_distributed.py

@@ -1516,7 +1516,9 @@ def _test_DistributedDataParallel_SyncBatchNorm(self, gpu_subset, rank, output_d
     def test_DistributedDataParallel_SyncBatchNorm(self):
         group, group_id, rank = self._init_global_test()
         rank_to_GPU = self._init_multigpu_helper()
-        gpus = list(rank_to_GPU[rank])
+        # DDP does not support replicating BN layers within a process, hence
+        # testing with one module replica per process
+        gpus = [rank]
         self._test_DistributedDataParallel_SyncBatchNorm(gpu_subset=gpus, rank=rank)
 
         # test output_device