Skip to content

Commit

Permalink
[Follow-up] Account for BC in PR#2947 (#2967)
Browse files Browse the repository at this point in the history
* Account for BC

* Remove an unnecessary stmt

---------

Co-authored-by: vfdev <vfdev.5@gmail.com>
  • Loading branch information
sadra-barikbin and vfdev-5 authored Jun 21, 2023
1 parent d85aacd commit dc1ddb6
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions ignite/distributed/comp_models/native.py
Original file line number Diff line number Diff line change
Expand Up @@ -436,14 +436,15 @@ def _do_all_gather(self, tensor: torch.Tensor, group: Optional[Any] = None) -> t
group_size = self.get_world_size()
elif isinstance(group, dist.ProcessGroup):
group_size = group.size()
elif isinstance(group, list):
group_size = len(group)
else:
raise ValueError("Argument group should be list of int or ProcessGroup")
if tensor.ndimension() == 0:
tensor = tensor.unsqueeze(0)
output = [torch.zeros_like(tensor) for _ in range(group_size)]
dist.all_gather(output, tensor, group=group)
if group is not None:
dist.all_gather(output, tensor, group=group)
else:
dist.all_gather(output, tensor)
return torch.cat(output, dim=0)

def _do_new_group(self, ranks: List[int], **kwargs: Any) -> Any:
Expand Down

0 comments on commit dc1ddb6

Please sign in to comment.