[Refactor] EPLB rebalance algo to NumPy (#30697)

Signed-off-by: ilmarkov <markovilya197@gmail.com>
2026-01-12 19:13:23 +01:00
parent 3d962d72ab
commit 1eb61ab34b
3 changed files with 126 additions and 128 deletions
--- a/tests/distributed/test_eplb_execute.py
+++ b/tests/distributed/test_eplb_execute.py
@@ -311,7 +311,7 @@ def _test_async_transfer_layer_without_mtp_worker(
            is_unchanged=is_unchanged,
            is_received_locally=is_received_locally,
            recv_metadata=recv_metadata,
-            new_indices=new_indices_cpu[layer_idx],
+            new_indices=new_indices_cpu[layer_idx].numpy(),
            ep_rank=ep_rank,
        )