def repeat_interleave_batch(x: Tensor, B: int, repeat: int) -> Tensor:
    """Repeat and interleave the input tensor.

    Args:
        x:
            Tensor with shape (B * N, ...) where B is the batch size and N the number of
            batches.
        B:
            Batch size.
        repeat:
            Number of times to repeat each batch.

    Returns:
        Tensor with shape (B * repeat * N, ...) where each batch is repeated `repeat`
        times.
    """
    N = len(x) // B
    x = torch.cat(
        [
            torch.cat([x[i * B : (i + 1) * B] for _ in range(repeat)], dim=0)
            for i in range(N)
        ],
        dim=0,
    )
    return x
Repeat and interleave the input tensor.

    Args:
        x:
            Tensor with shape (B * N, ...) where B is the batch size and N the number of
            batches.
        B:
            Batch size.
        repeat:
            Number of times to repeat each batch.

    Returns:
        Tensor with shape (B * repeat * N, ...) where each batch is repeated `repeat`
        times.
    
lightly/models/utils.py