def select_most_similar(
    x: Tensor,
    y: Tensor,
    y_values: Tensor,
) -> Tensor:
    """For each feature in x, searches the most similar feature in y and returns the
    corresponding value from y_values.

    Args:
        x:
            Tensor with shape (B, N, C).
        y:
            Tensor with shape (B, N, C).
        y_values:
            Tensor with shape (B, N, D).
    Returns:
        Values with shape (B, N, D) where values[i, j] is the entry in y_values[i, ...]
        such that x[i, j] is the most similar to y[i, ...].
    """
    y_index = most_similar_index(x, y)
    y_index = y_index.unsqueeze(-1).expand(y_values.shape)
    return torch.gather(y_values, dim=1, index=y_index)
For each feature in x, searches the most similar feature in y and returns the
    corresponding value from y_values.

    Args:
        x:
            Tensor with shape (B, N, C).
        y:
            Tensor with shape (B, N, C).
        y_values:
            Tensor with shape (B, N, D).
    Returns:
        Values with shape (B, N, D) where values[i, j] is the entry in y_values[i, ...]
        such that x[i, j] is the most similar to y[i, ...].
    
lightly/models/utils.py