Search before asking

• I have searched the YOLOv5 issues and found no similar bug report.

YOLOv5 Component

Export

Bug

Hello,
I saw issue #7268 which is pretty similar to mine. In the suggestion in that case, it mentions to not use --device but I need to use it to export to FP6 (--half). Thus, my error is

ONNX: export failure: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu! (when checking argument for argument index in method wrapper__index_select)

If I dont use --dynamic, then model is exported correctly, so my question is if its possible to get an FP16 model export with dynamic axes?. Thanks

Environment

• YOLO: YOLOv5 🚀 v6.1-165-g1a3ecb8 torch 1.10.0+cu113 CUDA:0 (NVIDIA GeForce RTX 3060 Laptop GPU, 5938MiB)
• OS: Ubuntu 20.04
• Python: 3.8.2

Minimal Reproducible Example

python export.py --weights yolov5s.pt --include onnx --half --device 0 --dynamic    

Additional

I would submit a PR, but I think this is a bit beyond my knowledge :)

Are you willing to submit a PR?

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