def decode(self, heatmaps: Tensor,
               displacements: Tensor) -> Tuple[np.ndarray, np.ndarray]:
        """Decode the keypoint coordinates from heatmaps and displacements. The
        decoded keypoint coordinates are in the input image space.

        Args:
            heatmaps (Tensor): Encoded root and keypoints (optional) heatmaps
                in shape (1, H, W) or (K+1, H, W)
            displacements (Tensor): Encoded keypoints displacement fields
                in shape (K*D, H, W)

        Returns:
            tuple:
            - keypoints (Tensor): Decoded keypoint coordinates in shape
                (N, K, D)
            - scores (tuple):
                - root_scores (Tensor): The root scores in shape (N, )
                - keypoint_scores (Tensor): The keypoint scores in
                    shape (N, K). If keypoint heatmaps are not generated,
                    `keypoint_scores` will be `None`
        """
        # heatmaps, displacements = encoded
        _k, h, w = displacements.shape
        k = _k // 2
        displacements = displacements.view(k, 2, h, w)

        # convert displacements to a dense keypoint prediction
        y, x = torch.meshgrid(torch.arange(h), torch.arange(w))
        regular_grid = torch.stack([x, y], dim=0).to(displacements)
        posemaps = (regular_grid[None] + displacements).flatten(2)

        # find local maximum on root heatmap
        root_heatmap_peaks = batch_heatmap_nms(heatmaps[None, -1:],
                                               self.decode_nms_kernel)
        root_scores, pos_idx = root_heatmap_peaks.flatten().topk(
            self.decode_max_instances)
        mask = root_scores > self.decode_thr
        root_scores, pos_idx = root_scores[mask], pos_idx[mask]

        keypoints = posemaps[:, :, pos_idx].permute(2, 0, 1).contiguous()

        if self.generate_keypoint_heatmaps and heatmaps.shape[0] == 1 + k:
            # compute scores for each keypoint
            keypoint_scores = self.get_keypoint_scores(heatmaps[:k], keypoints)
        else:
            keypoint_scores = None

        keypoints = torch.cat([
            kpt * self.scale_factor[i]
            for i, kpt in enumerate(keypoints.split(1, -1))
        ],
                              dim=-1)
        return keypoints, (root_scores, keypoint_scores)
Decode the keypoint coordinates from heatmaps and displacements. The
        decoded keypoint coordinates are in the input image space.

        Args:
            heatmaps (Tensor): Encoded root and keypoints (optional) heatmaps
                in shape (1, H, W) or (K+1, H, W)
            displacements (Tensor): Encoded keypoints displacement fields
                in shape (K*D, H, W)

        Returns:
            tuple:
            - keypoints (Tensor): Decoded keypoint coordinates in shape
                (N, K, D)
            - scores (tuple):
                - root_scores (Tensor): The root scores in shape (N, )
                - keypoint_scores (Tensor): The keypoint scores in
                    shape (N, K). If keypoint heatmaps are not generated,
                    `keypoint_scores` will be `None`
        
mmpose/mmpose/codecs/spr.py