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?