Input, output and indices must be on the current device解决

在linux环境，python3.8，pytorch1.8的环境，使用

torch.onnx._export将pytorch的.pt模型导出为.onnx的时候，之前一直可以正常导出的。

这次居然在这里报错，异常信息最后两行如下：

 params_dict = torch._C._jit_pass_onnx_constant_fold(graph, params_dict,
RuntimeError: Input, output and indices must be on the current device

异常提醒，输入或模型可能不在同一个设备上，但是我都检查过了。

input = input.to(device)

model = model.to(device)

也甚至怕model中间某些模块没有放进显卡，还用以下代码检查过了：

for param in model.parameters():

    assert param.device == device, "模型参数未移动到同一设备上"

都没有起作用。

最后解决是在，

关闭常量折叠优化： 常量折叠优化有时可能会引起设备不一致的问题，可以尝试关闭。

即在以下代码中do_constant_folding设置为False，即可关闭，本题异常也消失，解决。

torch.onnx.export(

model,

dummy_input,

'model.onnx',

export_params=True,

do_constant_folding=False)

顺便记录下，希望也能帮到一些小伙伴。