lightning module too

Author: awaelchli

src/lightning/pytorch/core/module.py

@@ -668,7 +668,8 @@ def all_gather(
 
         Return:
             A tensor of shape (world_size, batch, ...), or if the input was a collection
-            the output will also be a collection with tensors of this shape.
+            the output will also be a collection with tensors of this shape. For the special case where
+            world_size is 1, no additional dimension is added to the tensor(s).
 
         """
         group = group if group is not None else torch.distributed.group.WORLD