    def forward(self, image, question, answer=None, k=None, weights=None, train=True):
        """
        Args:
        k: number of answers for each question
        weights: weight for each answer
        """
        image_embeds = self.encode_vision(image)
        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(image.device)

        if train:
            answer_targets = answer.input_ids.masked_fill(
                answer.input_ids == self.tokenizer.pad_token_id, -100
            )

            question_output = self.text_encoder(
                question.input_ids,
                attention_mask=question.attention_mask,
                encoder_hidden_states=image_embeds,
                encoder_attention_mask=image_atts,
                return_dict=True,
            )

            question_states = []
            question_atts = []
            for b, n in enumerate(k):
                question_states += [question_output.last_hidden_state[b]] * n
                question_atts += [question.attention_mask[b]] * n
            question_states = torch.stack(question_states, 0)
            question_atts = torch.stack(question_atts, 0)

            answer_output = self.text_decoder(
                answer.input_ids,
                attention_mask=answer.attention_mask,
                encoder_hidden_states=question_states,
                encoder_attention_mask=question_atts,
                labels=answer_targets,
                return_dict=True,
                reduction="none",
            )
            loss = weights * answer_output.loss
            loss = loss.sum() / image.size(0)

            return loss

        # inference
        else:
            question_output = self.text_encoder(
                question.input_ids,
                attention_mask=question.attention_mask,
                encoder_hidden_states=image_embeds,
                encoder_attention_mask=image_atts,
                return_dict=True,
            )
            topk_ids, topk_probs = self.rank_answer(
                question_output.last_hidden_state,
                question.attention_mask,
                answer.input_ids,
                answer.attention_mask,
                k,
            )  # (bsz, 128), (bsz, 128)
            return topk_ids, topk_probs
Args:
k: number of answers for each question
weights: weight for each answer
videomamba/video_mm/models/umt_qa.py