mirror of
https://github.com/deepseek-ai/EPLB
synced 2025-05-05 04:24:26 +00:00
fix function return type
This commit is contained in:
commit
cf221658a7
2
eplb.py
2
eplb.py
@ -42,7 +42,7 @@ def balanced_packing(weight: torch.Tensor, num_packs: int) -> Tuple[torch.Tensor
|
||||
return pack_index, rank_in_pack
|
||||
|
||||
|
||||
def replicate_experts(weight: torch.Tensor, num_phy: int) -> torch.Tensor:
|
||||
def replicate_experts(weight: torch.Tensor, num_phy: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
|
||||
"""
|
||||
Replicate `num_log` experts to `num_phy` replicas, such that the maximum load of all replicas is minimized.
|
||||
|
Loading…
Reference in New Issue
Block a user