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

        Note:
            - batch_size: B
            - num_output_heatmap_levels: L
            - num_keypoints: K
            - heatmaps height: H
            - heatmaps weight: W
            - num_instances: N (usually 1 in topdown heatmap heads)

        Args:
            feats (Sequence[Sequence[Tensor]]): Feature maps from multiple
                stages and units
            batch_data_samples (List[:obj:`PoseDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance_labels` and `gt_fields`.
            train_cfg (Config, optional): The training config

        Returns:
            dict: A dictionary of loss components.
        """
        # multi-stage multi-unit predict heatmaps
        msmu_pred_heatmaps = self.forward(feats)

        keypoint_weights = torch.cat([
            d.gt_instance_labels.keypoint_weights for d in batch_data_samples
        ])  # shape: [B*N, L, K]

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

            # select `gt_heatmaps` and `keypoint_weights` for different level
            # according to `self.level_indices` to calculate loss
            gt_heatmaps = torch.stack([
                d.gt_fields[self.level_indices[i]].heatmaps
                for d in batch_data_samples
            ])
            loss_i = loss_func(msmu_pred_heatmaps[i], gt_heatmaps,
                               keypoint_weights[:, self.level_indices[i]])

            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(msmu_pred_heatmaps[-1]),
            target=to_numpy(gt_heatmaps),
            mask=to_numpy(keypoint_weights[:, -1]) > 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.

        Note:
            - batch_size: B
            - num_output_heatmap_levels: L
            - num_keypoints: K
            - heatmaps height: H
            - heatmaps weight: W
            - num_instances: N (usually 1 in topdown heatmap heads)

        Args:
            feats (Sequence[Sequence[Tensor]]): Feature maps from multiple
                stages and units
            batch_data_samples (List[:obj:`PoseDataSample`]): The Data
                Samples. It usually includes information such as
                `gt_instance_labels` and `gt_fields`.
            train_cfg (Config, optional): The training config

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