def _batch_shuffle(self, batch: Tensor) -> Tuple[Tensor, Tensor]:
        """Returns the shuffled batch and the indices to undo."""
        batch_size = batch.shape[0]
        shuffle = torch.randperm(batch_size, device=batch.device)
        return batch[shuffle], shuffle
Returns the shuffled batch and the indices to undo.
lightly/models/_momentum.py