[CI][Distributed] Fix edge case: One rank case (Rank 0) should get [False, False] (#146099)

nWEIdia · pytorchmergebot · commit 3a4e7a589b26 · 2025-01-31T20:31:13.000Z
To match the expected tensor (i.e. 2nd element in the array). Making rank0 receive [False, False] Fixes one of the issues reported in #146094 Pull Request resolved: #146099 Approved by: https://github.com/eqy
diff --git a/torch/testing/_internal/distributed/distributed_test.py b/torch/testing/_internal/distributed/distributed_test.py
@@ -6477,7 +6477,8 @@ def test_nccl_backend_bool_reduce(self):
             inp = {0: [True, True], 1: [False, False]}
             # Run reduce() with product op
             for op in [dist.ReduceOp.PRODUCT, dist.ReduceOp.MIN]:
-                input_tensor = torch.tensor(inp[self.rank % 2]).to(self.rank)
+                # make sure rank 0 gets False if WORLD_SIZE=1 to match expected tensor
+                input_tensor = torch.tensor(inp[(self.rank + 1) % 2]).to(self.rank)
                 expected = torch.tensor([False, False]).to(self.rank)
                 self._run_reduction_test(input_tensor, expected, op, dist.reduce, dst=0)
                 # Ensure that all ranks contributing True (cast to 1) results in the
