alltoall_v_inplace test

BaguaSys · Nov 9, 2021 · 5d5e2e4 · 5d5e2e4
1 parent b592958
commit 5d5e2e4
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/examples/communication_primitives/main.py b/examples/communication_primitives/main.py
@@ -165,11 +165,17 @@ def main():
         recv_sdispls,
         comm=comm,
     )
+    bagua.alltoall_v_inplace(send_tensors, send_counts, send_sdispls)
     assert torch.equal(
         recv_tensors, recv_tensor_bagua
     ), "recv_tensors:{a}, recv_tensor_bagua:{b}".format(
         a=recv_tensors, b=recv_tensor_bagua
     )
+    assert torch.equal(
+        send_tensors, recv_tensor_bagua
+    ), "recv_tensors:{a}, recv_tensor_bagua:{b}".format(
+        a=recv_tensors, b=recv_tensor_bagua
+    )
 
 
 if __name__ == "__main__":