def forward(self, pred: torch.Tensor, target: torch.Tensor,
                active_mask: torch.Tensor) -> torch.Tensor:
        """Forward function of MAE head.

        Args:
            pred (torch.Tensor): The reconstructed image.
            target (torch.Tensor): The target image.
            active_mask (torch.Tensor): The mask of the target image.
        Returns:
            torch.Tensor: The reconstruction loss.
        """
        # (B, C, H, W) -> (B, L, C) and perform normalization
        target = self.construct_target(target)

        # (B, C, H, W) -> (B, L, C)
        pred = self.patchify(pred)

        # (B, 1, f, f) -> (B, L)
        non_active_mask = active_mask.logical_not().int().view(
            active_mask.shape[0], -1)

        # MSE loss on masked patches
        loss = self.loss(pred, target, non_active_mask)
        return loss
Forward function of MAE head.

        Args:
            pred (torch.Tensor): The reconstructed image.
            target (torch.Tensor): The target image.
            active_mask (torch.Tensor): The mask of the target image.
        Returns:
            torch.Tensor: The reconstruction loss.
        
mmpretrain/models/heads/spark_head.py