def shard_xattn_Text_IMRAM(model, images, captions, caplens, iteration_step, shard_size=128):
            """
            Computer pairwise t2i image-caption distance with locality sharding
            """
            n_im_shard = int((len(images) - 1) / shard_size) + 1
            n_cap_shard = int((len(captions) - 1) / shard_size) + 1

            # print("n_im_shard: %d, n_cap_shard: %d" % (n_im_shard, n_cap_shard))

            d = [np.zeros((len(images), len(captions))) for _ in range(iteration_step)]

            for i in range(n_im_shard):
                im_start, im_end = shard_size * i, min(shard_size * (i + 1), len(images))
                for j in range(n_cap_shard):
                    cap_start, cap_end = shard_size * j, min(shard_size * (j + 1), len(captions))
                    im_emb = torch.from_numpy(images[im_start:im_end]).cuda()
                    s = torch.from_numpy(captions[cap_start:cap_end]).cuda()
                    l = caplens[cap_start:cap_end]
                    sim_list = model.xattn_score_Text_IMRAM(im_emb, s, l)
                    assert len(sim_list) == iteration_step
                    for k in range(iteration_step):
                        d[k][im_start:im_end, cap_start:cap_end] = sim_list[k].data.cpu().numpy()

            score = 0
            for j in range(iteration_step):
                score += d[j]

            return score

            Computer pairwise t2i image-caption distance with locality sharding
            
torchmm/models/retrieval/imram.py