diff --git a/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py b/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py index 6fcd90158f5e5cc028227803f0dd47bf6fd84dc1..6debf636958fe3a19f9485bd0cb8745a09bb3bcc 100644 --- a/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py +++ b/python/paddle/fluid/tests/unittests/collective/process_group_nccl.py @@ -568,6 +568,22 @@ class TestProcessGroupFp32(unittest.TestCase): print("test send api ok") + # test send 0-d tensor + # rank 0 + x = np.random.uniform(-1, 1, []).astype(self.dtype) + tensor_x = paddle.to_tensor(x) + # rank 1 + y = np.array(0.2022).astype(self.dtype) + tensor_y = paddle.to_tensor(y) + + if pg.rank() == 0: + task = dist.send(tensor_x, 1, sync_op=True) + else: + task = dist.recv(tensor_y, 0, sync_op=True) + assert np.array_equal(tensor_y, tensor_x) and tensor_y.shape == [] + + print("test send & recv 0-d tensor ok") + class TestProcessGroupFp16(TestProcessGroupFp32): def setUp(self):