code: def caffe2_keypoint_rcnn_inference(
      use_heatmap_max_keypoint, pred_keypoint_logits, pred_instances
  ):
      # just return the keypoint heatmap for now,
      # there will be option to call HeatmapMaxKeypointOp
      output = alias(pred_keypoint_logits, "kps_score")
      if all(isinstance(x, InstancesList) for x in pred_instances):
          assert len(pred_instances) == 1
          if use_heatmap_max_keypoint:
              device = output.device
              output = torch.ops._caffe2.HeatmapMaxKeypoint(
                  to_device(output, "cpu"),
                  pred_instances[0].pred_boxes.tensor,
                  should_output_softmax=True,  # worth make it configerable?
              )
              output = to_device(output, device)
              output = alias(output, "keypoints_out")
          pred_instances[0].set("pred_keypoints", output)
      return pred_keypoint_logits
  
imports: import math
  from typing import Dict
  import torch
  import torch.nn.functional
  from detectron2.layers import ShapeSpec, cat
  from detectron2.layers.roi_align_rotated import ROIAlignRotated
  from detectron2.modeling import poolers
  from detectron2.modeling.proposal_generator import rpn
  from detectron2.modeling.roi_heads.mask_head import mask_rcnn_inference
  from detectron2.structures import Boxes, ImageList, Instances, Keypoints, RotatedBoxes
  from .shared import alias, to_device
label_desc: Use image node_id 3 for calling a global function with example usage: caffe2_keypoint_rcnn_inference(use_heatmap_max_keypoint, pred_keypoint_logits, pred_instances) and returns: pred_keypoint_logits