def forward(
            self, logits: torch.Tensor, target: torch.Tensor,
            latent_pred: torch.Tensor,
            latent_target: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        """Forward function of CAE Loss.

        Args:
            logits (torch.Tensor): The outputs from the decoder.
            target (torch.Tensor): The targets generated by dalle.
            latent_pred (torch.Tensor): The latent prediction from the
                regressor.
            latent_target (torch.Tensor): The latent target from the teacher
                network.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: The main loss and align loss.
        """
        loss_main = self.loss_cross_entropy(logits, target)
        loss_align = self.loss_mse(latent_pred,
                                   latent_target.detach()) * self.lambd

        return loss_main, loss_align
Forward function of CAE Loss.

        Args:
            logits (torch.Tensor): The outputs from the decoder.
            target (torch.Tensor): The targets generated by dalle.
            latent_pred (torch.Tensor): The latent prediction from the
                regressor.
            latent_target (torch.Tensor): The latent target from the teacher
                network.

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: The main loss and align loss.
        
mmpretrain/models/losses/cae_loss.py