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_tags = self.forward(feats)

        if not self.tag_per_keypoint:
            pred_tags = pred_tags.repeat((1, self.num_keypoints, 1, 1))

        gt_heatmaps = torch.stack(
            [d.gt_fields.heatmaps for d in batch_data_samples])
        gt_masks = torch.stack(
            [d.gt_fields.heatmap_mask for d in batch_data_samples])
        keypoint_weights = torch.cat([
            d.gt_instance_labels.keypoint_weights for d in batch_data_samples
        ])
        keypoint_indices = [
            d.gt_instance_labels.keypoint_indices for d in batch_data_samples
        ]

        loss_kpt = self.loss_module.keypoint_loss(pred_heatmaps, gt_heatmaps,
                                                  keypoint_weights, gt_masks)

        loss_pull, loss_push = self.loss_module.tag_loss(
            pred_tags, keypoint_indices)

        losses = {
            'loss_kpt': loss_kpt,
            'loss_pull': loss_pull,
            'loss_push': loss_push
        }

        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/heatmap_heads/ae_head.py