Problem exporting to onnx, input output not on same device

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):
    def forward(self, x):
        return x[...,2:4]*self.ancor[0]

this example is possible to run on the cuda device:

x ='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/ 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)

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?

I don’t know, why this operation is failing, as I’m also unable to detect a device mismatch.
Would you mind creating an issue on GitHub so that we can track and check it?

Thanks for checking. Im currently on a trip and are not able to run the script to show the system setup. If you have time, please send in the repport. Otherwize, I will do it tomorrow.(bug report: Problem exporting to onnx with registered buffer · Issue #52432 · pytorch/pytorch · GitHub)