def METHOD_NAME(maps, rois):
    keypoints = torch.ops._caffe2.HeatmapMaxKeypoint(
        torch.tensor(maps),
        torch.tensor(rois),
        should_output_softmax=True,
    )
    return [keypoints.numpy()]