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

        Args:
            feats (Sequence[Tensor]): Multi-stage feature maps.
            batch_data_samples (List[:obj:`PoseDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instances`.
            train_cfg (Config, optional): The training config.

        Returns:
            dict: A dictionary of loss components.
        """
        multi_stage_pred_heatmaps = self.forward(feats)

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

        # calculate losses over multiple stages
        losses = dict()
        for i in range(self.num_stages):
            if isinstance(self.loss_module, nn.ModuleList):
                # use different loss_module over different stages
                loss_func = self.loss_module[i]
            else:
                # use the same loss_module over different stages
                loss_func = self.loss_module

            # the `gt_heatmaps` and `keypoint_weights` used to calculate loss
            # for different stages are the same
            loss_i = loss_func(multi_stage_pred_heatmaps[i], gt_heatmaps,
                               keypoint_weights)

            if 'loss_kpt' not in losses:
                losses['loss_kpt'] = loss_i
            else:
                losses['loss_kpt'] += loss_i

        # calculate accuracy
        _, avg_acc, _ = pose_pck_accuracy(
            output=to_numpy(multi_stage_pred_heatmaps[-1]),
            target=to_numpy(gt_heatmaps),
            mask=to_numpy(keypoint_weights) > 0)

        acc_pose = torch.tensor(avg_acc, device=gt_heatmaps.device)
        losses.update(acc_pose=acc_pose)

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

        Args:
            feats (Sequence[Tensor]): Multi-stage feature maps.
            batch_data_samples (List[:obj:`PoseDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instances`.
            train_cfg (Config, optional): The training config.

        Returns:
            dict: A dictionary of loss components.
        
mmpose/mmpose/models/heads/heatmap_heads/cpm_head.py