def invariance_loss(z: Tensor) -> Tensor:
    """Calculates the invariance loss, representing the similiarity between the patch embeddings and the average of
    the patch embeddings.

    Args:
        z:
            Patch embeddings.
    Returns:
        Similarity loss.
    """
    # z has shape (num_views, batch_size, dim)

    # Calculate the mean of the patch embeddings across the batch dimension
    z_mean = z.mean(0, keepdim=True)

    return -F.cosine_similarity(z, z_mean, dim=-1).mean()
Calculates the invariance loss, representing the similiarity between the patch embeddings and the average of
    the patch embeddings.

    Args:
        z:
            Patch embeddings.
    Returns:
        Similarity loss.
    
lightly/loss/emp_ssl_loss.py