Fix all_gather gloo test

Author: gruebel

ignite/distributed/comp_models/base.py

@@ -279,7 +279,9 @@ def all_reduce(self, tensor: Union[torch.Tensor, Number], op: str = "sum") -> Un
     def all_gather(
         self, tensor: Union[torch.Tensor, Number, str]
     ) -> Union[torch.Tensor, Number, List[Number], List[str]]:
-        return cast(Union[torch.Tensor, Number, List[Number], List[str]], tensor)
+        if isinstance(tensor, torch.Tensor):
+            return tensor
+        return cast(Union[List[Number], List[str]], [tensor])
 
     def broadcast(self, tensor: Union[torch.Tensor, Number, str], src: int = 0) -> Union[torch.Tensor, Number, str]:
         return tensor

tests/ignite/distributed/utils/__init__.py

@@ -83,7 +83,7 @@ def _test_distrib_all_reduce(device):
 
 def _test_distrib_all_gather(device):
 
-    res = idist.all_gather(10)
+    res = torch.tensor(idist.all_gather(10))
     true_res = torch.tensor([10,] * idist.get_world_size(), device=device)
     assert (res == true_res).all()
 

tests/ignite/distributed/utils/test_serial.py

@@ -54,4 +54,5 @@ def test_idist_all_reduce_no_dist():
 
 
 def test_idist_all_gather_no_dist():
-    assert idist.all_gather(10) == 10
+    assert idist.all_gather(10) == [10]
+    assert (idist.all_gather(torch.tensor(10)) == torch.tensor(10)).all()