def all_gather(
        self, data: Union[Tensor, dict, list, tuple], group: Optional[Any] = None, sync_grads: bool = False
    ) -> Union[Tensor, dict, list, tuple]:
        r"""Gather tensors or collections of tensors from multiple processes.

        This method needs to be called on all processes and the tensors need to have the same shape across all
        processes, otherwise your program will stall forever.

        Args:
            data: int, float, tensor of shape (batch, ...), or a (possibly nested) collection thereof.
            group: the process group to gather results from. Defaults to all processes (world)
            sync_grads: flag that allows users to synchronize gradients for the all_gather operation

        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. 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
        all_gather = self.trainer.strategy.all_gather
        data = convert_to_tensors(data, device=self.device)
        return apply_to_collection(data, Tensor, all_gather, group=group, sync_grads=sync_grads)
Gather tensors or collections of tensors from multiple processes.

        This method needs to be called on all processes and the tensors need to have the same shape across all
        processes, otherwise your program will stall forever.

        Args:
            data: int, float, tensor of shape (batch, ...), or a (possibly nested) collection thereof.
            group: the process group to gather results from. Defaults to all processes (world)
            sync_grads: flag that allows users to synchronize gradients for the all_gather operation

        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. For the special case where
            world_size is 1, no additional dimension is added to the tensor(s).

        
src/lightning/pytorch/core/module.py