In trainer_contrastive.py line129 feat = this_feat[:, perm[:K]] should be idxs_perm = idxs[perm] feat = this_feat[:,idxs_perm[:K]].squeeze() Otherwise, it can not select the current points.