def norm_mse_loss(x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
    """Normalized MSE Loss as implemented in https://github.com/htdt/self-supervised.

    Args:
        x0: First input tensor.
        x1: Second input tensor.

    Returns:
        The computed normalized MSE loss.
    """
    x0 = F.normalize(x0)
    x1 = F.normalize(x1)
    return torch.sub(input=2, other=(x0 * x1).sum(dim=-1).mean(), alpha=2)
Normalized MSE Loss as implemented in https://github.com/htdt/self-supervised.

    Args:
        x0: First input tensor.
        x1: Second input tensor.

    Returns:
        The computed normalized MSE loss.
    
lightly/loss/wmse_loss.py