def prototype_probabilities(
    queries: Tensor,
    prototypes: Tensor,
    temperature: float,
) -> Tensor:
    """Returns probability for each query to belong to each prototype.

    Args:
        queries:
            Tensor with shape (batch_size, dim)
        prototypes:
            Tensor with shape (num_prototypes, dim)
        temperature:
            Inverse scaling factor for the similarity.

    Returns:
        Probability tensor with shape (batch_size, num_prototypes) which sums to 1 along
        the num_prototypes dimension.
    """
    return F.softmax(torch.matmul(queries, prototypes.T) / temperature, dim=1)
Returns probability for each query to belong to each prototype.

    Args:
        queries:
            Tensor with shape (batch_size, dim)
        prototypes:
            Tensor with shape (num_prototypes, dim)
        temperature:
            Inverse scaling factor for the similarity.

    Returns:
        Probability tensor with shape (batch_size, num_prototypes) which sums to 1 along
        the num_prototypes dimension.
    
lightly/loss/msn_loss.py