update test_c10d.py and distributed_c10d.py

Signed-off-by: Jagadish Krishnamoorthy <jagdish.krishna@gmail.com>
pytorch · Dec 29, 2020 · d8ac57f · d8ac57f
1 parent a37fb05
commit d8ac57f
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/test/distributed/test_c10d.py b/test/distributed/test_c10d.py
@@ -4652,6 +4652,18 @@ def test_nccl_barrier_device_ids(self):
 
         c10d.barrier(device_ids=[self.rank])
 
+    @requires_nccl()
+    def test_nccl_barrier_device_ids_function_argument(self):
+        store = c10d.FileStore(self.file_name, self.world_size)
+        c10d.init_process_group(
+            backend="nccl",
+            rank=self.rank,
+            world_size=self.world_size,
+            store=store)
+
+        with self.assertRaisesRegex(RuntimeError, "Invalid function argument"):
+            c10d.barrier(device_ids=self.rank)
+
     @requires_gloo()
     def test_gloo_barrier_device_ids(self):
         store = c10d.FileStore(self.file_name, self.world_size)

diff --git a/torch/distributed/distributed_c10d.py b/torch/distributed/distributed_c10d.py
@@ -2386,7 +2386,7 @@ def barrier(group=GroupMember.WORLD,
         group (ProcessGroup, optional): The process group to work on. If None,
             the default process group will be used.
         async_op (bool, optional): Whether this op should be an async op
-        device_ids ([int], optional): List of device/GPU ids
+        device_ids ([int], optional): List of device/GPU ids.
                                       Valid only for NCCL backend.
 
     Returns: