def xattn_score_Image_IMRAM(self, images, captions_all, cap_lens):
        """
        Images: (batch_size, n_regions, d) matrix of images
        captions_all: (batch_size, max_n_words, d) matrix of captions
        CapLens: (batch_size) array of caption lengths
        """
        similarities = [[] for _ in range(self.iteration_step)]
        n_image = images.size(0)
        n_caption = captions_all.size(0)
        for i in range(n_caption):
            # Get the i-th text description
            n_word = cap_lens[i]
            cap_i = captions_all[i, :n_word, :].unsqueeze(0).contiguous()
            cap_i_expand = cap_i.repeat(n_image, 1, 1)

            query = images
            context = cap_i_expand
            for j in range(self.iteration_step):
                attn_feat, _ = func_attention(query, context, self.raw_feature_norm, smooth=self.lambda_softmax)

                row_sim = cosine_similarity(images, attn_feat, dim=2)
                row_sim = row_sim.mean(dim=1, keepdim=True)
                similarities[j].append(row_sim)

                query = self.gated_memory_i2t(query, attn_feat)

        # (n_image, n_caption)
        new_similarities = []
        for j in range(self.iteration_step):
            if len(similarities[j]) == 0:
                new_similarities.append([])
                continue
            similarities_one = torch.cat(similarities[j], 1)
            if self.training:
                similarities_one = similarities_one.transpose(0, 1)
            new_similarities.append(similarities_one)

        return new_similarities

        Images: (batch_size, n_regions, d) matrix of images
        captions_all: (batch_size, max_n_words, d) matrix of captions
        CapLens: (batch_size) array of caption lengths
        
torchmm/models/retrieval/imram.py