def predict(self,
                feats: Features,
                batch_data_samples: OptSampleList,
                test_cfg: ConfigType = {}) -> Predictions:
        """Predict results from features.

        Args:
            feats (Features): The features which could be in following forms:

                - Tuple[Tensor]: multi-stage features from the backbone
                - List[Tuple[Tensor]]: multiple features for TTA where either
                    `flip_test` or `multiscale_test` is applied
                - List[List[Tuple[Tensor]]]: multiple features for TTA where
                    both `flip_test` and `multiscale_test` are applied

            batch_data_samples (List[:obj:`PoseDataSample`]): The batch
                data samples
            test_cfg (dict): The runtime config for testing process. Defaults
                to {}

        Returns:
            Union[InstanceList | Tuple[InstanceList | PixelDataList]]: If
            ``test_cfg['output_heatmap']==True``, return both pose and heatmap
            prediction; otherwise only return the pose prediction.

            The pose prediction is a list of ``InstanceData``, each contains
            the following fields:

                - keypoints (np.ndarray): predicted keypoint coordinates in
                    shape (num_instances, K, D) where K is the keypoint number
                    and D is the keypoint dimension
                - keypoint_scores (np.ndarray): predicted keypoint scores in
                    shape (num_instances, K)

            The heatmap prediction is a list of ``PixelData``, each contains
            the following fields:

                - heatmaps (Tensor): The predicted heatmaps in shape (K, h, w)
        """
        # test configs
        multiscale_test = test_cfg.get('multiscale_test', False)
        flip_test = test_cfg.get('flip_test', False)
        shift_heatmap = test_cfg.get('shift_heatmap', False)
        align_corners = test_cfg.get('align_corners', False)
        restore_heatmap_size = test_cfg.get('restore_heatmap_size', False)
        output_heatmaps = test_cfg.get('output_heatmaps', False)

        # enable multi-scale test
        if multiscale_test:
            # TTA: multi-scale test
            assert is_list_of(feats, list if flip_test else tuple)
        else:
            assert is_list_of(feats, tuple if flip_test else Tensor)
            feats = [feats]

        # resize heatmaps to align with with input size
        if restore_heatmap_size:
            img_shape = batch_data_samples[0].metainfo['img_shape']
            assert all(d.metainfo['img_shape'] == img_shape
                       for d in batch_data_samples)
            img_h, img_w = img_shape
            heatmap_size = (img_w, img_h)
        else:
            heatmap_size = None

        multiscale_heatmaps = []
        multiscale_tags = []

        for scale_idx, _feats in enumerate(feats):
            if not flip_test:
                _heatmaps, _tags = self.forward(_feats)

            else:
                # TTA: flip test
                assert isinstance(_feats, list) and len(_feats) == 2
                flip_indices = batch_data_samples[0].metainfo['flip_indices']
                # original
                _feats_orig, _feats_flip = _feats
                _heatmaps_orig, _tags_orig = self.forward(_feats_orig)

                # flipped
                _heatmaps_flip, _tags_flip = self.forward(_feats_flip)
                _heatmaps_flip = flip_heatmaps(
                    _heatmaps_flip,
                    flip_mode='heatmap',
                    flip_indices=flip_indices,
                    shift_heatmap=shift_heatmap)
                _tags_flip = self._flip_tags(
                    _tags_flip,
                    flip_indices=flip_indices,
                    shift_heatmap=shift_heatmap)

                # aggregated heatmaps
                _heatmaps = aggregate_heatmaps(
                    [_heatmaps_orig, _heatmaps_flip],
                    size=heatmap_size,
                    align_corners=align_corners,
                    mode='average')

                # aggregated tags (only at original scale)
                if scale_idx == 0:
                    _tags = aggregate_heatmaps([_tags_orig, _tags_flip],
                                               size=heatmap_size,
                                               align_corners=align_corners,
                                               mode='concat')
                else:
                    _tags = None

            multiscale_heatmaps.append(_heatmaps)
            multiscale_tags.append(_tags)

        # aggregate multi-scale heatmaps
        if len(feats) > 1:
            batch_heatmaps = aggregate_heatmaps(
                multiscale_heatmaps,
                align_corners=align_corners,
                mode='average')
        else:
            batch_heatmaps = multiscale_heatmaps[0]
        # only keep tags at original scale
        batch_tags = multiscale_tags[0]

        batch_outputs = tuple([batch_heatmaps, batch_tags])
        preds = self.decode(batch_outputs)

        if output_heatmaps:
            pred_fields = []
            for _heatmaps, _tags in zip(batch_heatmaps.detach(),
                                        batch_tags.detach()):
                pred_fields.append(PixelData(heatmaps=_heatmaps, tags=_tags))

            return preds, pred_fields
        else:
            return preds
Predict results from features.

        Args:
            feats (Features): The features which could be in following forms:

                - Tuple[Tensor]: multi-stage features from the backbone
                - List[Tuple[Tensor]]: multiple features for TTA where either
                    `flip_test` or `multiscale_test` is applied
                - List[List[Tuple[Tensor]]]: multiple features for TTA where
                    both `flip_test` and `multiscale_test` are applied

            batch_data_samples (List[:obj:`PoseDataSample`]): The batch
                data samples
            test_cfg (dict): The runtime config for testing process. Defaults
                to {}

        Returns:
            Union[InstanceList | Tuple[InstanceList | PixelDataList]]: If
            ``test_cfg['output_heatmap']==True``, return both pose and heatmap
            prediction; otherwise only return the pose prediction.

            The pose prediction is a list of ``InstanceData``, each contains
            the following fields:

                - keypoints (np.ndarray): predicted keypoint coordinates in
                    shape (num_instances, K, D) where K is the keypoint number
                    and D is the keypoint dimension
                - keypoint_scores (np.ndarray): predicted keypoint scores in
                    shape (num_instances, K)

            The heatmap prediction is a list of ``PixelData``, each contains
            the following fields:

                - heatmaps (Tensor): The predicted heatmaps in shape (K, h, w)
        
mmpose/mmpose/models/heads/heatmap_heads/ae_head.py