def loss(self,
             feats: Tuple[Tensor],
             batch_data_samples: OptSampleList,
             train_cfg: ConfigType = {}) -> dict:
        """Calculate losses from a batch of inputs and data samples.

        Args:
            feats (Tuple[Tensor]): The multi-stage features
            batch_data_samples (List[:obj:`PoseDataSample`]): The batch
                data samples
            train_cfg (dict): The runtime config for training process.
                Defaults to {}

        Returns:
            dict: A dictionary of losses.
        """
        pred_heatmaps, pred_displacements = self.forward(feats)
        gt_heatmaps = torch.stack(
            [d.gt_fields.heatmaps for d in batch_data_samples])
        heatmap_weights = torch.stack(
            [d.gt_fields.heatmap_weights for d in batch_data_samples])
        gt_displacements = torch.stack(
            [d.gt_fields.displacements for d in batch_data_samples])
        displacement_weights = torch.stack(
            [d.gt_fields.displacement_weights for d in batch_data_samples])

        if 'heatmap_mask' in batch_data_samples[0].gt_fields.keys():
            heatmap_mask = torch.stack(
                [d.gt_fields.heatmap_mask for d in batch_data_samples])
        else:
            heatmap_mask = None

        # calculate losses
        losses = dict()
        heatmap_loss = self.loss_module['heatmap'](pred_heatmaps, gt_heatmaps,
                                                   heatmap_weights,
                                                   heatmap_mask)
        displacement_loss = self.loss_module['displacement'](
            pred_displacements, gt_displacements, displacement_weights)

        losses.update({
            'loss/heatmap': heatmap_loss,
            'loss/displacement': displacement_loss,
        })

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

        Args:
            feats (Tuple[Tensor]): The multi-stage features
            batch_data_samples (List[:obj:`PoseDataSample`]): The batch
                data samples
            train_cfg (dict): The runtime config for training process.
                Defaults to {}

        Returns:
            dict: A dictionary of losses.
        
mmpose/mmpose/models/heads/hybrid_heads/dekr_head.py