    def vtm_loss(
        self,
        multimodal_encoder,
        vtm_head: nn.Module,
        temp,
        vision_embeds: torch.Tensor,
        text_embeds: torch.Tensor,
        vision_proj: torch.Tensor,
        text_proj: torch.Tensor,
        text_atts: torch.Tensor,
        idx: torch.Tensor,
    ):
        """video-text matching loss.

        Args:
            multinomial_encoder (nn.Module): The multimodal_encoder.
            vtm_head (nn.Module): The head to produce the video-text matching score.
            temp (torch.Tensor): temporature for similarity calculation.
            vision_embeds (torch.Tensor): The features of all patches in the video. Shape: [B,T,L,C].
            text_embeds (torch.Tensor): The features of all tokens in the text. Shape: [B,L,C].
            vision_proj (torch.Tensor): The vision representation. Shape: [B,T,C].
            text_proj (torch.Tensor): The text representation. Shape: [B,C].
            text_atts (torch.Tensor): The padded mask for text tokens. 0 is padded. Shape: [B,L].
            idx (torch.Tensor): The index for each example. Shape: [B,].

        Returns: TODO

        """
        with torch.no_grad():
            sim_v2t, sim_t2v = get_sim(vision_proj, text_proj, temp)
            vision_atts = torch.ones(
                vision_embeds.size()[:-1], dtype=torch.long, device=vision_embeds.device
            )
            weights_v2t = F.softmax(sim_v2t + 1e-4, dim=1)  # (N, N)
            weights_t2v = F.softmax(sim_t2v + 1e-4, dim=1)

            mask = self.get_mask(sim_v2t, idx=idx).bool()
            weights_v2t.masked_fill_(mask, 0)
            weights_t2v.masked_fill_(mask, 0)
            weights_v2t = torch.nan_to_num_(weights_v2t, nan=1e-2, posinf=1e-2, neginf=1e-2)
            weights_t2v = torch.nan_to_num_(weights_t2v, nan=1e-2, posinf=1e-2, neginf=1e-2)

        # select a negative image for each text
        if self.vtm_hard_neg:
            vision_neg_indices = torch.multinomial(weights_t2v, 1).squeeze()
            txt_neg_indices = torch.multinomial(weights_v2t, 1).squeeze()
        else:
            vision_neg_indices = self.get_rand_indices(mask, 1).squeeze()
            txt_neg_indices = self.get_rand_indices(mask, 1).squeeze()

        vision_embeds_neg = vision_embeds[vision_neg_indices]  # [B, T*L, c]
        text_embeds_neg = text_embeds[txt_neg_indices]  # [B, L, d]
        text_atts_neg = text_atts[txt_neg_indices]

        # concat embeddings
        vision_embeds_all = torch.cat([vision_embeds, vision_embeds_neg, vision_embeds], dim=0)
        text_embeds_all = torch.cat([text_embeds, text_embeds, text_embeds_neg], dim=0)
        vision_atts_all = torch.cat([vision_atts, vision_atts, vision_atts], dim=0)
        text_atts_all = torch.cat([text_atts, text_atts, text_atts_neg], dim=0)

        output = multimodal_encoder(
            encoder_embeds=text_embeds_all,
            attention_mask=text_atts_all,
            encoder_hidden_states=vision_embeds_all,
            encoder_attention_mask=vision_atts_all,
            return_dict=True,
            mode="fusion",
        )

        vtm_embeds = output.last_hidden_state[:, 0]  # pos (N, d) + neg (2N, d)

        vtm_logits = vtm_head(vtm_embeds)  # [3*B, 2]

        bs = vtm_logits.shape[0] // 3
        vtm_labels = vtm_logits.new_ones(3 * bs, dtype=torch.long)
        vtm_labels[bs:] = 0
        loss_vtm = F.cross_entropy(vtm_logits, vtm_labels)
        return loss_vtm
video-text matching loss.

Args:
multinomial_encoder (nn.Module): The multimodal_encoder.
vtm_head (nn.Module): The head to produce the video-text matching score.
temp (torch.Tensor): temporature for similarity calculation.
vision_embeds (torch.Tensor): The features of all patches in the video. Shape: [B,T,L,C].
text_embeds (torch.Tensor): The features of all tokens in the text. Shape: [B,L,C].
vision_proj (torch.Tensor): The vision representation. Shape: [B,T,C].
text_proj (torch.Tensor): The text representation. Shape: [B,C].
text_atts (torch.Tensor): The padded mask for text tokens. 0 is padded. Shape: [B,L].
idx (torch.Tensor): The index for each example. Shape: [B,].

Returns: TODO
videomamba/video_mm/models/criterions.py