[Minor] Code style improvements (#2355)

This commit is contained in:
Lianmin Zheng
2024-12-04 19:02:08 -08:00
committed by GitHub
parent 9cc733b38c
commit 2b0fc5941d
3 changed files with 16 additions and 16 deletions

View File

@@ -47,7 +47,7 @@ def _to_torch(model: torch.nn.Module, reverse: bool, batch_size: int):
if "FusedMoE" in sub.__class__.__name__:
if batch_size == 1:
# The performance of torch.compile on this layer is not always good when bs > 1,
# so we decide to skip it for now.
# so we decide to only use torch.compile when bs =1
sub._forward_method = fused_moe_forward_native
else:
sub._forward_method = sub.forward_native