IndexError: too many indices for tensor of dimension 3 - pytorch

Question

I am trying to convert Detectron2 Model into an onnx model format (pytorch to onnx). This is the code I am using (note I am using all necessary imports and to my understanding my installations are correct). Here is my code:
cfg = get_cfg()
cfg.MODEL.DEVICE = 'cpu'
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"))
model = build_model(cfg)
aug = T.ResizeShortestEdge([cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST],
cfg.INPUT.MAX_SIZE_TEST)
height, width = im.shape[:2]
image = aug.get_transform(im).apply_image(im)
image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
inputs2 = {"image": image}
print(image.shape)
output = export_onnx_model(model, image)
print("output:", output)
onnx.save(output, "/home/ecoation/Documents/model/deploy.onnx")
And here is the bug I am getting:
Traceback (most recent call last):
File "/home/ecoation/Documents/detectronCode.py", line 57, in <module>
output = export_onnx_model(model, image)
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/detectron2/export/caffe2_export.py", line 49, in export_onnx_model
torch.onnx.export(
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/onnx/__init__.py", line 272, in export
return utils.export(model, args, f, export_params, verbose, training,
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/onnx/utils.py", line 88, in export
_export(model, args, f, export_params, verbose, training, input_names, output_names,
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/onnx/utils.py", line 694, in _export
_model_to_graph(model, args, verbose, input_names,
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/onnx/utils.py", line 457, in _model_to_graph
graph, params, torch_out, module = _create_jit_graph(model, args,
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/onnx/utils.py", line 420, in _create_jit_graph
graph, torch_out = _trace_and_get_graph_from_model(model, args)
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/jit/_trace.py", line 116, in wrapper
outs.append(self.inner(*trace_inputs))
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 887, in _call_impl
result = self._slow_forward(*input, **kwargs)
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/torch/nn/modules/module.py", line 860, in _slow_forward
result = self.forward(*input, **kwargs)
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 150, in forward
return self.inference(batched_inputs)
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 203, in inference
images = self.preprocess_image(batched_inputs)
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 228, in preprocess_image
images = [self._move_to_current_device(x["image"]) for x in batched_inputs]
File "/home/ecoation/anaconda3/envs/detectron2/lib/python3.8/site-packages/detectron2/modeling/meta_arch/rcnn.py", line 228, in <listcomp>
images = [self._move_to_current_device(x["image"]) for x in batched_inputs]
IndexError: too many indices for tensor of dimension 2
Feel free to take a look, any help in turning this into a onnx format model is much appreciated. If it helps, I am on Ubuntu 18.04 and all I'm trying to do is take a pretrained simple detectron2 model from regular to onnx format. I'm hoping to get it so I can do custom, but any help with pretrained would be much appreciated.
Edit:
Image size is this :
torch.Size([3, 800, 1067])