yolov3_onnx sample with non-square input resolution

I have a YOLOv3 trained on custom object which works well.
Since i found the sample to use it on TensorRT i want to give a try to see if i can improve the performance on my TX2.
I have reshape problem and i think that’s because my input resolution is (HW) 640x832

File "onnx_to_tensorrt.py", line 169, in main
    boxes, classes, scores = postprocessor.process(trt_outputs, shape_orig_WH)
  File "/home/god_kane/Documents/python_projects/yolov3_onnx/data_processing.py", line 126, in process
    boxes, categories, confidences = self._process_yolo_output(outputs_reshaped, resolution_raw)
  File "/home/god_kane/Documents/python_projects/yolov3_onnx/data_processing.py", line 161, in _process_yolo_output
    box, category, confidence = self._process_feats(output, mask)
  File "/home/god_kane/Documents/python_projects/yolov3_onnx/data_processing.py", line 240, in _process_feats
    col = col.reshape(grid_h, grid_w, 1, 1).repeat(3, axis=-2)
ValueError: cannot reshape array of size 676 into shape (20,26,1,1)

What i changed is :

In data_processing.py i made sure i loaded my classes.txt files and the assert checks 3

In yolov3_to_onnx.py :

output_tensor_dims['082_convolutional'] = [24, 20, 26]
    output_tensor_dims['094_convolutional'] = [24, 40, 52]
    output_tensor_dims['106_convolutional'] = [24, 80, 104]

I have 3 classes so C is 24. For the output dimension this is a wild guess based on the 608 input.
It may not be the right value but i don"t know where to find that information.

In onnx_to_tensorrt.py :

output_shapes = [(1, 24, 20, 26), (1, 24, 40, 52), (1, 24, 80, 104)]

Again here, wild guess.

The onnx file and trt file are created successfully but interpreting the result fails with the exception above.
An size of 676 is obviously (26,26,1,1) but since my input is not a square it doesn’t work.
I am not familiar at all about the code in _process_feats and how i should interpret the inference result.
Since the inference runs, it’s only a problem in building the result and print the box on the original image.

Thanks a lot for your insight and help.