Skip to content

Commit

Permalink
Updated typehint for idist.all_gather method (#3089)
Browse files Browse the repository at this point in the history
Description:
- This reflects the fact that idist can use all_gather_object internally
  • Loading branch information
vfdev-5 committed Oct 11, 2023
1 parent f26078f commit 854b5ea
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 6 deletions.
8 changes: 4 additions & 4 deletions ignite/distributed/comp_models/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,7 @@ def all_reduce(

def all_gather(
self, tensor: Union[torch.Tensor, float, str, Any], group: Optional[Any] = None
) -> Union[torch.Tensor, float, List[float], List[str]]:
) -> Union[torch.Tensor, float, List[float], List[str], List[Any]]:
if not isinstance(tensor, (torch.Tensor, Number, str)):
return self._do_all_gather_object(tensor, group=group)

Expand Down Expand Up @@ -355,11 +355,11 @@ def all_reduce(
return tensor

def all_gather(
self, tensor: Union[torch.Tensor, float, str], group: Optional[Any] = None
) -> Union[torch.Tensor, float, List[float], List[str]]:
self, tensor: Union[torch.Tensor, float, str, Any], group: Optional[Any] = None
) -> Union[torch.Tensor, float, List[float], List[str], List[Any]]:
if isinstance(tensor, torch.Tensor):
return tensor
return cast(Union[List[float], List[str]], [tensor])
return cast(Union[List[float], List[str], List[Any]], [tensor])

def broadcast(
self, tensor: Union[torch.Tensor, float, str, None], src: int = 0, safe_mode: bool = False
Expand Down
4 changes: 2 additions & 2 deletions ignite/distributed/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -351,8 +351,8 @@ def all_reduce(


def all_gather(
tensor: Union[torch.Tensor, float, str], group: Optional[Union[Any, List[int]]] = None
) -> Union[torch.Tensor, float, List[float], List[str]]:
tensor: Union[torch.Tensor, float, str, Any], group: Optional[Union[Any, List[int]]] = None
) -> Union[torch.Tensor, float, List[float], List[str], List[Any]]:
"""Helper method to perform all gather operation.
Args:
Expand Down

0 comments on commit 854b5ea

Please sign in to comment.