fix format

microsoft · Mar 4, 2024 · 55dbaf5 · 55dbaf5
1 parent eb33fb1
commit 55dbaf5
Showing 1 changed file with 1 addition and 4 deletions.
diff --git a/deepspeed/moe/sharded_moe.py b/deepspeed/moe/sharded_moe.py
@@ -95,10 +95,7 @@ def gumbel_rsample(shape: Tuple, device: torch.device) -> Tensor:
 class _AllToAll(torch.autograd.Function):
 
  @staticmethod
- def forward(
- ctx: Any,
- group: dist.ProcessGroup,
- input: Tensor) -> Tensor: # type: ignore
+ def forward(ctx: Any, group: dist.ProcessGroup, input: Tensor) -> Tensor: # type: ignore
  ctx.group = group
  input = input.contiguous()
  output = torch.empty_like(input)