import torch


def get_closest_features_to_simplex(simplex, crops_features):
    """
    Return closest image crops to each of the simplex summets
    Args:
        simplex: simplex features
        crops_features: image crops
    Returns:
        closest_crops: closest image crops to each of the simplex summets
    """
    dim = crops_features.shape[-1]
    distances = torch.norm(simplex.reshape(-1, 1, dim) - crops_features.reshape(1, -1, dim), dim=2, p=2)
    closest_crops = torch.argmin(distances, dim=1)
    return torch.stack([crops_features[index_aug] for index_aug in closest_crops])