def batch_unshuffle_distributed(
    batch: torch.Tensor, shuffle: torch.Tensor
) -> torch.Tensor:
    """Undo batch shuffle over multiple devices.

    This code was taken and adapted from here:
    https://github.com/facebookresearch/moco.

    Args:
        batch:
            The tensor to unshuffle.
        shuffle:
            Index to restore the original tensor.

    Returns:
        The unshuffled tensor.

    """
    # gather from all devices
    batch_size_this = batch.shape[0]
    batch_gather = concat_all_gather(batch)
    batch_size_all = batch_gather.shape[0]

    # Calculate the number of devices
    num_devices = batch_size_all // batch_size_this

    # Get the rank of the current device
    rank = dist.get_rank()

    # Index for this device after unshuffling
    idx_this = shuffle.view(num_devices, -1)[rank]

    # Returns the unshuffled batch for this device
    return batch_gather[idx_this]
Undo batch shuffle over multiple devices.

    This code was taken and adapted from here:
    https://github.com/facebookresearch/moco.

    Args:
        batch:
            The tensor to unshuffle.
        shuffle:
            Index to restore the original tensor.

    Returns:
        The unshuffled tensor.

    
lightly/models/utils.py