def loss(self, inputs: Tensor, data_samples: SampleList) -> dict:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            inputs (Tensor): Inputs with shape (N, C, H, W).
            data_samples (List[:obj:`PoseDataSample`]): The batch
                data samples.

        Returns:
            dict: A dictionary of losses.
        """
        self.set_epoch()

        losses = dict()
        
        with torch.no_grad():
            fea_t = self.teacher.extract_feat(inputs)
            lt_x, lt_y = self.teacher.head(fea_t)
            pred_t = (lt_x, lt_y)

        if not self.two_dis:
            fea_s = self.student.extract_feat(inputs)
            ori_loss, pred, gt, target_weight = self.head_loss(fea_s, data_samples, train_cfg=self.train_cfg)
            losses.update(ori_loss)
        else:
            ori_loss, pred, gt, target_weight = self.head_loss(fea_t, data_samples, train_cfg=self.train_cfg)

        all_keys = self.distill_losses.keys()

        if 'loss_fea' in all_keys:
            loss_name = 'loss_fea'
            losses[loss_name] = self.distill_losses[loss_name](fea_s[-1], fea_t[-1])
            if not self.two_dis:
                losses[loss_name] = (1-self.epoch/self.max_epochs)*losses[loss_name]

        if 'loss_logit' in all_keys:
            loss_name = 'loss_logit'
            losses[loss_name] = self.distill_losses[loss_name](pred, pred_t, self.student.head.loss_module.beta, target_weight)
            if not self.two_dis:
                losses[loss_name] = (1-self.epoch/self.max_epochs)*losses[loss_name]

        return losses
Calculate losses from a batch of inputs and data samples.

        Args:
            inputs (Tensor): Inputs with shape (N, C, H, W).
            data_samples (List[:obj:`PoseDataSample`]): The batch
                data samples.

        Returns:
            dict: A dictionary of losses.
        
mmpose/mmpose/models/pose_estimators/distiller.py