From e116a89020b058760da84b11be4a124ebabf3b64 Mon Sep 17 00:00:00 2001
From: youkaichao <youkaichao@gmail.com>
Date: Mon, 2 Dec 2024 16:52:57 -0800
Subject: [PATCH] improve tests

Signed-off-by: youkaichao <youkaichao@gmail.com>
---
 tests/distributed/test_pynccl.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py
index 6f1498a5d0852..b7390759d72f8 100644
--- a/tests/distributed/test_pynccl.py
+++ b/tests/distributed/test_pynccl.py
@@ -319,8 +319,7 @@ def broadcast_worker_fn():
 
     for i in range(pynccl_comm.world_size):
         pynccl_comm.broadcast(recv_tensors[i], src=i)
-        result = recv_tensors[i].mean().cpu().item()
-        assert result == i
+        assert torch.all(recv_tensors[i] == i).cpu().item()
 
 
 def test_ncclGetUniqueId():