def loss(self,
             pred: torch.Tensor,
             target: torch.Tensor,
             mask: Optional[torch.Tensor] = None) -> torch.Tensor:
        """Forward head.

        Args:
            pred (torch.Tensor): Predictions with shape B x L x C.
            target (torch.Tensor): Targets with shape B x L x C.
            mask (torch.Tensor): Mask with shape B x L.

        Returns:
            torch.Tensor: The loss tensor.
        """
        loss = self.loss_module(pred, target, mask)
        return loss
Forward head.

        Args:
            pred (torch.Tensor): Predictions with shape B x L x C.
            target (torch.Tensor): Targets with shape B x L x C.
            mask (torch.Tensor): Mask with shape B x L.

        Returns:
            torch.Tensor: The loss tensor.
        
mmpretrain/models/heads/mim_head.py