ONNXRuntimeError for fasterrcnn_mobilenet_v3_large_fpn

🐛 Describe the bug

I get an ONNXRuntimeError when running FasterRCNN model exported to .onnx

Here is a minimal reproductible example:

obs: The error shows up only in the cropped image. If I comment out img = img[50:200, 150:250] - it works fine.

import torch
import torchvision
import cv2
import requests
import numpy as np
import onnxruntime

print("Exporting model")
model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights='DEFAULT')
torch.onnx.export(model, torch.rand(1, 3, 640, 640), '/tmp/model.onnx', 
                    input_names=['input'], output_names=['boxes', 'scores', 'labels'])


print("Downloading image")
r = requests.get('https://docs.opencv.org/4.x/roi.jpg', allow_redirects=True)
open('/tmp/roi.jpg', 'wb').write(r.content)
img = cv2.imread('/tmp/roi.jpg')

img = img[50:200, 150:250]
cv2.imwrite('/tmp/roi2.jpg', img)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
image_dim_dim = cv2.resize(img, (640, 640))
image_dim_dim = np.array(image_dim_dim, dtype=np.float32) / 255.0
image_bchw = np.transpose(np.expand_dims(image_dim_dim, 0), (0, 3, 1, 2))

print("Running inference")
session = onnxruntime.InferenceSession('/tmp/model.onnx', providers=["CPUExecutionProvider"])
outputs = [o.name for o in session.get_outputs()]
inputs = [o.name for o in session.get_inputs()]
prediction = session.run(outputs, {inputs[0]: image_bchw})
print(prediction)

Error:
onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException: [ONNXRuntimeError] : 6 : RUNTIME_EXCEPTION : Non-zero status code returned while running Reshape node. Name:'/roi_heads/Reshape_2' Status Message: /onnxruntime_src/onnxruntime/core/providers/cpu/tensor/reshape_helper.h:39 onnxruntime::ReshapeHelper::ReshapeHelper(const onnxruntime::TensorShape&, onnxruntime::TensorShapeVector&, bool) size != 0 && (input_shape_size % size) == 0 was false. The input tensor cannot be reshaped to the requested shape. Input shape:{490,363}, requested shape:{-1,4}

Versions

onnxruntime 1.18.1
torch 2.2.2
torchvision 0.17.2

Run the minimal example with the listed PyTorch, torchvision, and ONNXRuntime versions, comparing the full image with the cropped input. Trace the exported model's /roi_heads/Reshape_2 failure during session.run; done means the cropped image runs without the reshape error and the regression is covered by an appropriate test.