def _sample_feats(self, feats: Tensor, indices: Tensor) -> Tensor:
        """Extract feature vectors at the specified indices from the input
        feature map.

        Args:
            feats (Tensor): Input feature map.
            indices (Tensor): Indices of the feature vectors to extract.

        Returns:
            Tensor: Extracted feature vectors.
        """
        assert indices.dtype == torch.long
        if indices.shape[1] == 3:
            b, w, h = [ind.squeeze(-1) for ind in indices.split(1, -1)]
            instance_feats = feats[b, :, h, w]
        elif indices.shape[1] == 2:
            w, h = [ind.squeeze(-1) for ind in indices.split(1, -1)]
            instance_feats = feats[:, :, h, w]
            instance_feats = instance_feats.permute(0, 2, 1)
            instance_feats = instance_feats.reshape(-1,
                                                    instance_feats.shape[-1])

        else:
            raise ValueError(f'`indices` should have 2 or 3 channels, '
                             f'but got f{indices.shape[1]}')
        return instance_feats
Extract feature vectors at the specified indices from the input
        feature map.

        Args:
            feats (Tensor): Input feature map.
            indices (Tensor): Indices of the feature vectors to extract.

        Returns:
            Tensor: Extracted feature vectors.
        
mmpose/mmpose/models/heads/heatmap_heads/cid_head.py