Skip to content

Commit d215cef

Browse files
author
offline0806
committed
[EPLB]Add type check for local_count.
Signed-off-by: offline0806 <z00858301@china.huawei.com>
1 parent 2c3bd98 commit d215cef

File tree

1 file changed

+3
-2
lines changed

1 file changed

+3
-2
lines changed

vllm_ascend/eplb/core/eplb_utils.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,8 +47,9 @@ def determine_default_expert_map(global_expert_num, world_size, rank_id,
4747
else:
4848
start -= 1
4949

50-
local_ids = torch.arange(local_count, dtype=torch.int32)
51-
expert_map[start:end] = local_ids
50+
if isinstance(local_count, int):
51+
local_ids = torch.arange(local_count, dtype=torch.int32)
52+
expert_map[start:end] = local_ids
5253

5354
return (local_count, expert_map)
5455

0 commit comments

Comments
 (0)