See this discussion: https://github.com/KevinMusgrave/pytorch-metric-learning/discussions/561#discussioncomment-4580712