未验证 提交 6727ab5c 编写于 作者: W wuhuachaocoding 提交者: GitHub

update 0-d tensor test for send_recv. (#49218)

上级 68b78f30
...@@ -568,6 +568,22 @@ class TestProcessGroupFp32(unittest.TestCase): ...@@ -568,6 +568,22 @@ class TestProcessGroupFp32(unittest.TestCase):
print("test send api ok") 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): class TestProcessGroupFp16(TestProcessGroupFp32):
def setUp(self): def setUp(self):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册