def most_similar_index(
    x: Tensor,
    y: Tensor,
) -> Tensor:
    """For each feature in x, searches the most similar feature in y and returns the
    corresponding index.

    Args:
        x:
            Tensor with shape (B, N, C) containing the features to compare.
        y:
            Tensor with shape (B, N, C) containing the features to search for similarity.
    Returns:
        Index with shape (B, N) such that y[i, index[i, j]] is most similar to x[i, j]
        over all y[i, ...].

    """
    # Normalize the input tensors along the last dimension
    x = functional.normalize(x, dim=-1)
    y = functional.normalize(y, dim=-1)

    similarity = torch.einsum("bnc,bmc->bnm", x, y)
    return similarity.argmax(dim=2)
For each feature in x, searches the most similar feature in y and returns the
    corresponding index.

    Args:
        x:
            Tensor with shape (B, N, C) containing the features to compare.
        y:
            Tensor with shape (B, N, C) containing the features to search for similarity.
    Returns:
        Index with shape (B, N) such that y[i, index[i, j]] is most similar to x[i, j]
        over all y[i, ...].

    
lightly/models/utils.py