def loss(self, inputs: torch.Tensor, data_samples: List[DataSample],
             **kwargs) -> Dict[str, torch.Tensor]:
        """The forward function in training.

        Args:
            inputs (torch.Tensor): The input images.
            data_samples (List[DataSample]): All elements required
                during the forward function.

        Returns:
            Dict[str, torch.Tensor]: A dictionary of loss components.
        """
        # ids_restore: the same as that in original repo, which is used
        # to recover the original order of tokens in decoder.
        clip_feature, importance = self.target_generator(inputs)
        importance = importance[:, 0, 1:]
        latent, ids_restore, ids_keep, ids_dump = self.backbone(
            inputs, importance)
        pred = self.neck(latent, ids_restore, ids_keep, ids_dump)

        loss = self.head.loss(pred, clip_feature)
        losses = dict(loss=loss)
        return losses
The forward function in training.

        Args:
            inputs (torch.Tensor): The input images.
            data_samples (List[DataSample]): All elements required
                during the forward function.

        Returns:
            Dict[str, torch.Tensor]: A dictionary of loss components.
        
mmpretrain/models/selfsup/milan.py