Exporting native Keypoint R-CNN to ONNX using export_model.py fails.

I stumbled upon this issue when trying to convert a custom trained Mask R-CNN model with attached Keypoint head using the R50-DC5 backbone to onnx format.
At first I thought my model is the issue but I was able to reproduce the conversion error using the provided pretrained `keypoint_rcnn_R_50_FPN_3x` model.
Using the `mask_rcnn_R_50_FPN_3x` or `mask_rcnn_R_50_DC5_3x` the conversion worked flawless. It seems to be an issue with the Keypoint R-CNN architecture.

I am willing to contribute to fixing this issue. But I need more information on what exactly went wrong.

## Instructions To Reproduce the Issue:
I used the Nvidia PyTorch Docker however issue persists in other environments.
```
docker run --name nvidiatestenv --runtime=nvidia --gpus all -it nvcr.io/nvidia/pytorch:23.10-py3
```
Install Detectron2.
```
git clone https://github.com/facebookresearch/detectron2.git
pip install -e ./detectron2
```

Get the provided Keypoint-RCNN weights.
`wget https://dl.fbaipublicfiles.com/detectron2/COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x/137849621/model_final_a6e10b.pkl`

Make changes to `export_model.py` in lines 160-162 as stated by [NVIDIA TensorRT](https://github.com/NVIDIA/TensorRT/blob/3aaa97b91ee1dd61ea46f78683d9a3438f26192e/samples/python/detectron2/README.md)
```
aug = T.ResizeShortestEdge(
[cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST
)
```
has to be changed to:
```
aug = T.ResizeShortestEdge(
[1344, 1344], 1344
)
```

Download [sample image](https://github.com/facebookresearch/detectron2/assets/77785743/fc2d1ed9-0abe-422a-923b-5655497b77d1) and run export_model.py
```
curl -o 1344x1344.jpg https://github-production-user-asset-6210df.s3.amazonaws.com/77785743/281084531-fc2d1ed9-0abe-422a-923b-5655497b77d1.jpg
python detectron2/tools/deploy/export_model.py \
--sample-image 1344x1344.jpg \
--config-file detectron2/configs/COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml \
--export-method tracing \
--format onnx \
--output ./ \
MODEL.WEIGHTS ./model_final_a6e10b.pkl \
MODEL.DEVICE cuda
```
## Expected behavior:

Produce `model.onnx`

## Observed behaviour
```
Traceback (most recent call last):
File "/workspace/detectron2/tools/deploy/export_model.py", line 225, in
exported_model = export_tracing(torch_model, sample_inputs)
File "/workspace/detectron2/tools/deploy/export_model.py", line 132, in export_tracing
torch.onnx.export(traceable_model, (image,), f, opset_version=STABLE_ONNX_OPSET_VERSION)
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 516, in export
_export(
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1582, in _export
graph, params_dict, torch_out = _model_to_graph(
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1139, in _model_to_graph
graph = _optimize_graph(
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 677, in _optimize_graph
graph = _C._jit_pass_onnx(graph, operator_export_type)
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1926, in _run_symbolic_function
return symbolic_fn(graph_context, *inputs, **attrs)
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset9.py", line 6976, in prim_loop
torch._C._jit_pass_onnx_block(
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/utils.py", line 1926, in _run_symbolic_function
return symbolic_fn(graph_context, *inputs, **attrs)
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset11.py", line 999, in index
return opset9.index(g, self, index)
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_opset9.py", line 5847, in index
return symbolic_helper._unimplemented(
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_helper.py", line 612, in _unimplemented
_onnx_unsupported(f"{op}, {msg}", value)
File "/usr/local/lib/python3.10/dist-packages/torch/onnx/symbolic_helper.py", line 623, in _onnx_unsupported
raise errors.SymbolicValueError(
torch.onnx.errors.SymbolicValueError: Unsupported: ONNX export of operator aten::index, operator of advanced indexing on tensor of unknown rank. Try turning on shape inference during export: torch.onnx._export(..., onnx_shape_inference=True).. Please feel free to request support or submit a pull request on PyTorch GitHub: https://github.com/pytorch/pytorch/issues [Caused by the value 'roi_map.3 defined in (%roi_map.3 : Tensor = onnx::Reshape(%roi_map, %2701) # /workspace/detectron2/detectron2/structures/keypoints.py:205:18
)' (type 'Tensor') in the TorchScript graph. The containing node has kind 'onnx::Reshape'.]
(node defined in File "/workspace/detectron2/detectron2/structures/keypoints.py", line 205
# Although semantically equivalent, `reshape` is used instead of `squeeze` due
# to limitation during ONNX export of `squeeze` in scripting mode
roi_map = roi_map.reshape(roi_map.shape[1:]) # keypoints x H x W
~~~~~~~~~~~~~~~ <--- HERE

# softmax over the spatial region
)

Inputs:
#0: roi_map defined in (%roi_map : Tensor = onnx::Resize[coordinate_transformation_mode="half_pixel", cubic_coeff_a=-0.75, mode="cubic", nearest_mode="floor"](%2684, %2693, %2694, %2692) # /workspace/detectron2/detectron2/structures/keypoints.py:201:18
) (type 'Tensor')
#1: 2701 defined in (%2701 : LongTensor(device=cpu)[] = onnx::Slice(%2696, %2698, %2699, %2697, %2700) # /workspace/detectron2/detectron2/structures/keypoints.py:205:34
) (type 'List[Tensor]')
Outputs:
#0: roi_map.3 defined in (%roi_map.3 : Tensor = onnx::Reshape(%roi_map, %2701) # /workspace/detectron2/detectron2/structures/keypoints.py:205:18
) (type 'Tensor')
```

## What I tried
I tried suggestions from #4834 where a modification of the `export_model.py` script is suggested using the PR #4205.
However the exact same error appears.

## Further Information
Conversion to torchscript using `--export-method tracing/scripting` works.
However converting the resulting `model.ts` to TensorRT using `torch_tensorrt` produces a completely different error.
Try it out:
```
import torch
import torchvision
import torch_tensorrt
img = torchvision.io.read_image("1344x1344.jpg")
model = torch.jit.load("model.ts")
trt_ts_module = torch_tensorrt.compile(model,
inputs = [torch_tensorrt.Input(shape=img.shape, dtype=torch.float)],
enabled_precisions= {torch.float, torch.half}, ir='torchscript', truncate_long_and_double=True)

```

Error when using traced `model.ts`:
```
RuntimeError: [Error thrown at core/partitioning/shape_analysis.cpp:183] Expected ivalues_maps.count(input) to be true but got false
Could not find torch::jit::Value* width_corrections.1 produced from %width_corrections.1 : Tensor = aten::div(%widths.1, %widths_ceil.1) # /workspace/detectron2/detectron2/structures/keypoints.py:194:24 in lowering graph for mini graph input.
```
Error when using scripted `model.ts`:
`[ff42823c1538:462 :0:462] Caught signal 11 (Segmentation fault: address not mapped to object at address 0x60)
`
## Environment:

Paste the output of the following command:
```
------------------------------- --------------------------------------------------------------
sys.platform linux
Python 3.10.12 (main, Jun 11 2023, 05:26:28) [GCC 11.4.0]
numpy 1.22.2
detectron2 0.6 @/workspace/detectron2/detectron2
Compiler GCC 11.4
CUDA compiler CUDA 12.2
detectron2 arch flags 5.2, 6.0, 6.1, 7.0, 7.2, 7.5, 8.0, 8.6, 8.7, 9.0
DETECTRON2_ENV_MODULE
PyTorch 2.1.0a0+32f93b1 @/usr/local/lib/python3.10/dist-packages/torch
PyTorch debug build False
torch._C._GLIBCXX_USE_CXX11_ABI True
GPU available Yes
GPU 0 NVIDIA GeForce MX330 (arch=6.1)
Driver version 545.84
CUDA_HOME /usr/local/cuda
TORCH_CUDA_ARCH_LIST 5.2 6.0 6.1 7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX
Pillow 9.2.0
torchvision 0.16.0a0 @/usr/local/lib/python3.10/dist-packages/torchvision
torchvision arch flags 5.2, 6.0, 6.1, 7.0, 7.2, 7.5, 8.0, 8.6, 8.7, 9.0
fvcore 0.1.5.post20221221
iopath 0.1.9
cv2 4.7.0
------------------------------- --------------------------------------------------------------
PyTorch built with:
- GCC 11.2
- C++ Version: 201703
- Intel(R) oneAPI Math Kernel Library Version 2021.1-Product Build 20201104 for Intel(R) 64 architecture applications
- Intel(R) MKL-DNN v2.7.3 (Git Hash N/A)
- OpenMP 201511 (a.k.a. OpenMP 4.5)
- LAPACK is enabled (usually provided by MKL)
- NNPACK is enabled
- CPU capability usage: AVX2
- CUDA Runtime 12.2
- NVCC architecture flags: -gencode;arch=compute_52,code=sm_52;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_72,code=sm_72;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_80,code=sm_80;-gencode;arch=compute_86,code=sm_86;-gencode;arch=compute_87,code=sm_87;-gencode;arch=compute_90,code=sm_90;-gencode;arch=compute_90,code=compute_90
- CuDNN 8.9.5
- Magma 2.6.2
- Build settings: BLAS_INFO=mkl, BUILD_TYPE=Release, CUDA_VERSION=12.2, CUDNN_VERSION=8.9.5, CXX_COMPILER=/opt/rh/gcc-toolset-11/root/usr/bin/c++, CXX_FLAGS=-fno-gnu-unique -D_GLIBCXX_USE_CXX11_ABI=1 -fvisibility-inlines-hidden -DUSE_PTHREADPOOL -DNDEBUG -DUSE_KINETO -DLIBKINETO_NOROCTRACER -DUSE_FBGEMM -DUSE_QNNPACK -DUSE_PYTORCH_QNNPACK -DUSE_XNNPACK -DSYMBOLICATE_MOBILE_DEBUG_HANDLE -O2 -fPIC -Wall -Wextra -Werror=return-type -Werror=non-virtual-dtor -Werror=range-loop-construct -Werror=bool-operation -Wnarrowing -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-unused-parameter -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-stringop-overflow -Wno-psabi -Wno-error=pedantic -Wno-error=old-style-cast -Wno-invalid-partial-specialization -Wno-unused-private-field -Wno-aligned-allocation-unavailable -Wno-missing-braces -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Werror=format -Werror=cast-function-type -Wno-stringop-overflow, LAPACK_INFO=mkl, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, TORCH_DISABLE_GPU_ASSERTS=ON, TORCH_VERSION=2.1.0, USE_CUDA=ON, USE_CUDNN=ON, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=OFF, USE_NCCL=1, USE_NNPACK=ON, USE_OPENMP=ON, USE_ROCM=OFF,
```

Start by reproducing the command in tools/deploy/export_model.py, then inspect the failing path in structures/keypoints.py around lines 194-205. Compare tracing behavior for the keypoint and mask models; done means the provided keypoint_rcnn_R_50_FPN_3x model exports a usable model.onnx without the advanced-indexing error.

Written by the indexing model from the issue text.