def negative_mises_fisher_weights(
    out0: Tensor, out1: Tensor, sigma: float = 0.5
) -> Tensor:
    """Negative Mises-Fisher weighting function as presented in Decoupled Contrastive Learning [0].

    The implementation was inspired by [1].

    - [0] Chun-Hsiao Y. et. al., 2021, Decoupled Contrastive Learning https://arxiv.org/abs/2110.06848
    - [1] https://github.com/raminnakhli/Decoupled-Contrastive-Learning

    Args:
        out0:
            Output projections of the first set of transformed images.
            Shape: (batch_size, embedding_size)
        out1:
            Output projections of the second set of transformed images.
            Shape: (batch_size, embedding_size)
        sigma:
            Similarities are scaled by inverse sigma.
    Returns:
        A tensor with shape (batch_size,) where each entry is the weight for one
        of the input images.
    """
    similarity = torch.einsum("nm,nm->n", out0.detach(), out1.detach()) / sigma

    # Return negative Mises-Fisher weights
    return torch.tensor(2 - out0.shape[0] * nn.functional.softmax(similarity, dim=0))
Negative Mises-Fisher weighting function as presented in Decoupled Contrastive Learning [0].

    The implementation was inspired by [1].

    - [0] Chun-Hsiao Y. et. al., 2021, Decoupled Contrastive Learning https://arxiv.org/abs/2110.06848
    - [1] https://github.com/raminnakhli/Decoupled-Contrastive-Learning

    Args:
        out0:
            Output projections of the first set of transformed images.
            Shape: (batch_size, embedding_size)
        out1:
            Output projections of the second set of transformed images.
            Shape: (batch_size, embedding_size)
        sigma:
            Similarities are scaled by inverse sigma.
    Returns:
        A tensor with shape (batch_size,) where each entry is the weight for one
        of the input images.
    
lightly/loss/dcl_loss.py