As mentioned in this pull request: https://github.com/KevinMusgrave/pytorch-metric-learning/pull/424#issuecomment-1042868787