Met exception while exporting to onnx

[sylyt62]

I've already setup the model and inputs as:

device(type='cuda', index=0)

model = attempt_load(args.weights, map_location=device) # load FP32 model
im = torch.zeros(1, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection

But still met the following exeption:

Traceback (most recent call last):
File "D:\sdk\envs\smpl-py38-torch110-cu111\lib\site-packages\torch\onnx\utils.py", line 1529, in _export
graph, params_dict, torch_out = _model_to_graph(
File "D:\sdk\envs\smpl-py38-torch110-cu111\lib\site-packages\torch\onnx\utils.py", line 1172, in _model_to_graph
params_dict = _C._jit_pass_onnx_constant_fold(
RuntimeError: 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)
python-BaseException