def METHOD_NAME(self, key_embeds, ref_embeds, key_sampling_results,
          ref_sampling_results):
    num_key_rois = [res.pos_masks.size(0) for res in key_sampling_results]
    key_embeds = torch.split(key_embeds, num_key_rois)
    num_ref_rois = [res.pos_masks.size(0) for res in ref_sampling_results]
    ref_embeds = torch.split(ref_embeds, num_ref_rois)
    dists, cos_dists = [], []
    for key_embed, ref_embed in zip(key_embeds, ref_embeds):
        dist = cal_similarity(
            key_embed,
            ref_embed,
            method='dot_product',
            temperature=self.softmax_temp)
        dists.append(dist)
        if self.loss_track_aux is not None:
            cos_dist = cal_similarity(
                key_embed, ref_embed, method='cosine')
            cos_dists.append(cos_dist)
        else:
            cos_dists.append(None)
    return dists, cos_dists