BUG: YOLOv6 medium, large and X MBLA models do not torchscript export

[josht000]

Search before asking

• I have searched the YOLOv6 issues and found no similar feature requests.

Description

YOLOv6 medium, large and x do not torchscript export using any version of the export_torchscript.py files.

Example minimal output for the large model:

Sizes of tensors must match except in dimension 1. Expected size 160 but got size 4 for tensor number 1 in the list. File "/raid/home/thamesj/yolov6/YOLOv6-main/yolov6/models/effidehead.py", line 122, in return tuple(torch.cat([cls, reg], 1) for cls, reg in zip(cls_score_list, reg_dist_list)) File "/raid/home/thamesj/yolov6/YOLOv6-main/yolov6/models/effidehead.py", line 122, in forward return tuple(torch.cat([cls, reg], 1) for cls, reg in zip(cls_score_list, reg_dist_list)) File "/raid/home/thamesj/yolov6/YOLOv6-main/yolov6/models/yolo.py", line 40, in forward x = self.detect(x) File "/raid/home/thamesj/yolov6/YOLOv6-main/deploy/Torchscript/export_torchscript.py", line 59, in y = model(img) # dry run RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 160 but got size 4 for tensor number 1 in the list.

ChatGPT was able to get me through the error, but it cropped up again during inference, so I'm putting it back to you all to provide expert guidance.

Thanks for all you've done.

Use case

No response

Additional

No response

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!