    def vtc_loss(
        self,
        vision_proj: torch.Tensor,
        text_proj: torch.Tensor,
        idx: torch.Tensor,
        temp=1.0,
        all_gather=True,
    ):
        """forward to calculate the loss

        Args:
            vision_proj (torch.Tensor): The vision representation. Shape: [B,T,C].
            text_proj (torch.Tensor): The text representation. Shape: [B,C].
            idx (torch.Tensor): The index for each example. Shape: [B,].
            temp (torch.Tensor): The temperature. Shape: [].
            all_gather (bool): If true, will gather samples across all the GPUs and calculate loss across the gathered samples.

        Returns: loss_vtc (torch.Tensor): The video-text contrastive loss. Shape: [].

        """
        if all_gather:
            gather_args = self.get_gather_args()
            vision_proj = allgather_wgrad(vision_proj, gather_args)
            text_proj = allgather_wgrad(text_proj, gather_args)
            if idx is not None:
                idx = allgather_wgrad(idx, gather_args)

        sim_v2t, sim_t2v = get_sim(vision_proj, text_proj, temp)

        with torch.no_grad():
            sim_v2t_targets = self.get_mask(sim_v2t, idx=idx, normalize=True)
            sim_t2v_targets = sim_v2t_targets

        loss_i2t = -torch.sum(F.log_softmax(sim_v2t, dim=1) * sim_v2t_targets, dim=1).mean()
        loss_t2i = -torch.sum(F.log_softmax(sim_t2v, dim=1) * sim_t2v_targets, dim=1).mean()

        loss_vtc = (loss_i2t + loss_t2i) / 2
        return loss_vtc
forward to calculate the loss

Args:
vision_proj (torch.Tensor): The vision representation. Shape: [B,T,C].
text_proj (torch.Tensor): The text representation. Shape: [B,C].
idx (torch.Tensor): The index for each example. Shape: [B,].
temp (torch.Tensor): The temperature. Shape: [].
all_gather (bool): If true, will gather samples across all the GPUs and calculate loss across the gathered samples.

Returns: loss_vtc (torch.Tensor): The video-text contrastive loss. Shape: [].
videomamba/video_mm/models/criterions.py