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.
        """

        # load targets
        gt_heatmaps, gt_instance_coords, keypoint_weights = [], [], []
        heatmap_mask = []
        instance_imgids, gt_instance_heatmaps = [], []
        for i, d in enumerate(batch_data_samples):
            gt_heatmaps.append(d.gt_fields.heatmaps)
            gt_instance_coords.append(d.gt_instance_labels.instance_coords)
            keypoint_weights.append(d.gt_instance_labels.keypoint_weights)
            instance_imgids.append(
                torch.ones(
                    len(d.gt_instance_labels.instance_coords),
                    dtype=torch.long) * i)

            instance_heatmaps = d.gt_fields.instance_heatmaps.reshape(
                -1, self.num_keypoints,
                *d.gt_fields.instance_heatmaps.shape[1:])
            gt_instance_heatmaps.append(instance_heatmaps)

            if 'heatmap_mask' in d.gt_fields:
                heatmap_mask.append(d.gt_fields.heatmap_mask)

        gt_heatmaps = torch.stack(gt_heatmaps)
        heatmap_mask = torch.stack(heatmap_mask) if heatmap_mask else None

        gt_instance_coords = torch.cat(gt_instance_coords, dim=0)
        gt_instance_heatmaps = torch.cat(gt_instance_heatmaps, dim=0)
        keypoint_weights = torch.cat(keypoint_weights, dim=0)
        instance_imgids = torch.cat(instance_imgids).to(gt_heatmaps.device)

        # feed-forward
        feats = feats[-1]
        pred_instance_feats, pred_heatmaps = self.iia_module.forward_train(
            feats, gt_instance_coords, instance_imgids)

        # conpute contrastive loss
        contrastive_loss = 0
        for i in range(len(batch_data_samples)):
            pred_instance_feat = pred_instance_feats[instance_imgids == i]
            contrastive_loss += self.loss_module['contrastive'](
                pred_instance_feat)
        contrastive_loss = contrastive_loss / max(1, len(instance_imgids))

        # limit the number of instances
        max_train_instances = train_cfg.get('max_train_instances', -1)
        if (max_train_instances > 0
                and len(instance_imgids) > max_train_instances):
            selected_indices = torch.randperm(
                len(instance_imgids),
                device=gt_heatmaps.device,
                dtype=torch.long)[:max_train_instances]
            gt_instance_coords = gt_instance_coords[selected_indices]
            keypoint_weights = keypoint_weights[selected_indices]
            gt_instance_heatmaps = gt_instance_heatmaps[selected_indices]
            instance_imgids = instance_imgids[selected_indices]
            pred_instance_feats = pred_instance_feats[selected_indices]

        # calculate the decoupled heatmaps for each instance
        pred_instance_heatmaps = self.gfd_module(feats, pred_instance_feats,
                                                 gt_instance_coords,
                                                 instance_imgids)

        # calculate losses
        losses = {
            'loss/heatmap_coupled':
            self.loss_module['heatmap_coupled'](pred_heatmaps, gt_heatmaps,
                                                None, heatmap_mask)
        }
        if len(instance_imgids) > 0:
            losses.update({
                'loss/heatmap_decoupled':
                self.loss_module['heatmap_decoupled'](pred_instance_heatmaps,
                                                      gt_instance_heatmaps,
                                                      keypoint_weights),
                'loss/contrastive':
                contrastive_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/heatmap_heads/cid_head.py