I’m working with yolov5 and trying to export it to onnx but so far without any success.
So I have tried to narrow down the problem into a small code example to illustrate the problem.
torch version =1.7.1
x = torch.rand([2, 3, 60, 80, 6]).to('cuda:0')
ag = torch.tensor([[[[[[ 10., 13.]]],
[[[ 16., 30.]]],
[[[ 33., 23.]]]]],
[[[[[ 30., 61.]]],
[[[ 62., 45.]]],
[[[ 59., 119.]]]]],
[[[[[116., 90.]]],
[[[156., 198.]]],
[[[373., 326.]]]]]])
class Test(nn.Module):
def __init__(self, ancor):
super().__init__()
self.register_buffer('ancor',ancor)
def forward(self, x):
return x[...,2:4]*self.ancor[0]
this example is possible to run on the cuda device:
model=Test(ag).to('cuda:0').eval()
x = x.to('cuda:0')
result = model(x)
but when I try to export it then it becomes a problem:
fname = 'test.onnx'
torch.onnx.export(model, x, fname, verbose=False, opset_version=12)
part of the error message:
/opt/conda/lib/python3.7/site-packages/torch/onnx/utils.py in _model_to_graph(model, args, verbose, input_names, output_names, operator_export_type, example_outputs, _retain_param_name, do_constant_folding, _disable_torch_constant_prop, fixed_batch_size, training, use_new_jit_passes, dynamic_axes)
448 if do_constant_folding and _export_onnx_opset_version in torch.onnx.constant_folding_opset_versions:
449 params_dict = torch._C._jit_pass_onnx_constant_fold(graph, params_dict,
→ 450 _export_onnx_opset_version)
451 torch._C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph)
452
RuntimeError: Input, output and indices must be on the current device
I do not understand the error message, both tensors are on the same device.
Is this a bug or how do I avoid this error?