ONNX: export failure

[ZhangLe-fighting]

Before submitting a bug report, please be aware that your issue must be reproducible with all of the following, otherwise it is non-actionable, and we can not help you:

• Current repo: run git fetch && git status -uno to check and git pull to update repo
• Common dataset: coco.yaml or coco128.yaml
• Common environment: Colab, Google Cloud, or Docker image. See https://github.com/ultralytics/yolov5#environments

If this is a custom dataset/training question you must include your train*.jpg, val*.jpg and results.png figures, or we can not help you. You can generate these with utils.plot_results().

🐛 Bug

A clear and concise description of what the bug is.

To Reproduce (REQUIRED)

Input:

python models/export.py --weights yolov5_coco_model/yolov5_coco.pt  --img 640 --batch 1 --simplify --export-nms --device='0'

Output:

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)

Expected behavior

A clear and concise description of what you expected to happen.

Environment

If applicable, add screenshots to help explain your problem.

• CUDA:11.2
• GPU [Tesla P40]
• Linux hik-gpu 4.15.0-142-generic #146~16.04.1-Ubuntu SMP Tue Apr 13 09:27:15 UTC 2021 x86_64 x86_64 x86_64 GNU/Linux
• python3.8

Additional context

just copy the zip file and run export.py， if the input no --device='0'，success ！

[fratou]

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) i have same question

[K-tang-mkv]

change --device 0 to --device cpu