Torch.onnx.export, RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

username: cuixing158_1
cooked: <p>when i put model and input tensor convert to cuda device, then export onnx, occur above errors"RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!",my model is complicated, how to solve?</p>

username: ptrblck
cooked: <p>Is your model generally working via <code>torch.jit.script</code> or <code>torch.jit.trace</code> or are you seeing the same issue?<br>
Could you post an executable code snippet to reproduce the issue, please?</p>

username: Jordan_Chipka
cooked: <p>I’m encountering a similar issue. I’m attempting to convert a pytorch model to onnx with fp16 precision. I’m using the following command:</p>
<pre><code class="lang-auto">    torch.onnx.export(
        model,
        input_tensor,
        onnx_file_path,
        input_names=["input"],
        output_names=["output"],
        export_params=True,
    )
</code></pre>
<p>Both <code>model</code> and <code>input_tensor</code> are fp16 and on gpu (<code>model.cuda()</code>, <code>model.half()</code>, etc.). But I still get the following error:</p>
<pre><code class="lang-auto">File "/home/{USER}/.local/lib/python3.8/site-packages/torch/onnx/utils.py", line 628, in _model_to_graph
    params_dict = torch._C._jit_pass_onnx_eval_peephole(graph, params_dict)
RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!
</code></pre>
<p>When I inspect <code>graph</code> in line 628 of <code>site-packages/torch/onnx/utils.py</code>, I see that there appears to be some onnx::Constant tensors (among a few others) that did not make it onto gpu. See the full graph below (notice lines %39, %40, %44, %45, %50, %51, %55, %56, %61, %62, %66, %67, %72, %73, %77, %78, %79, %80, %81).</p>
<pre><code class="lang-auto">graph(%input : Half(1, 3, 32, 32, strides=[3, 1, 96, 3], requires_grad=0, device=cuda:0),
      %features.0.weight : Half(32, 3, 3, 3, strides=[27, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.0.bias : Half(32, strides=[1], requires_grad=1, device=cuda:0),
      %features.1.running_mean : Half(32, strides=[1], requires_grad=0, device=cuda:0),
      %features.1.running_var : Half(32, strides=[1], requires_grad=0, device=cuda:0),
      %features.1.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %features.4.weight : Half(32, 32, 3, 3, strides=[288, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.4.bias : Half(32, strides=[1], requires_grad=1, device=cuda:0),
      %features.5.running_mean : Half(32, strides=[1], requires_grad=0, device=cuda:0),
      %features.5.running_var : Half(32, strides=[1], requires_grad=0, device=cuda:0),
      %features.5.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %features.9.weight : Half(64, 32, 3, 3, strides=[288, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.9.bias : Half(64, strides=[1], requires_grad=1, device=cuda:0),
      %features.10.running_mean : Half(64, strides=[1], requires_grad=0, device=cuda:0),
      %features.10.running_var : Half(64, strides=[1], requires_grad=0, device=cuda:0),
      %features.10.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %features.13.weight : Half(64, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.13.bias : Half(64, strides=[1], requires_grad=1, device=cuda:0),
      %features.14.running_mean : Half(64, strides=[1], requires_grad=0, device=cuda:0),
      %features.14.running_var : Half(64, strides=[1], requires_grad=0, device=cuda:0),
      %features.14.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %features.18.weight : Half(128, 64, 3, 3, strides=[576, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.18.bias : Half(128, strides=[1], requires_grad=1, device=cuda:0),
      %features.19.running_mean : Half(128, strides=[1], requires_grad=0, device=cuda:0),
      %features.19.running_var : Half(128, strides=[1], requires_grad=0, device=cuda:0),
      %features.19.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %features.22.weight : Half(128, 128, 3, 3, strides=[1152, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.22.bias : Half(128, strides=[1], requires_grad=1, device=cuda:0),
      %features.23.running_mean : Half(128, strides=[1], requires_grad=0, device=cuda:0),
      %features.23.running_var : Half(128, strides=[1], requires_grad=0, device=cuda:0),
      %features.23.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %features.27.weight : Half(256, 128, 3, 3, strides=[1152, 9, 3, 1], requires_grad=1, device=cuda:0),
      %features.27.bias : Half(256, strides=[1], requires_grad=1, device=cuda:0),
      %features.28.running_mean : Half(256, strides=[1], requires_grad=0, device=cuda:0),
      %features.28.running_var : Half(256, strides=[1], requires_grad=0, device=cuda:0),
      %features.28.num_batches_tracked : Long(requires_grad=0, device=cuda:0),
      %classifier.0.weight : Half(10, 256, strides=[256, 1], requires_grad=1, device=cuda:0),
      %classifier.0.bias : Half(10, strides=[1], requires_grad=1, device=cuda:0)):
  %38 : Half(1, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%input, %features.0.weight, %features.0.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %39 : Half(32, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %40 : Half(32, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %41 : Half(1, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%38, %39, %40, %features.1.running_mean, %features.1.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %42 : Half(1, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%41) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %43 : Half(1, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%42, %features.4.weight, %features.4.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %44 : Half(32, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %45 : Half(32, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %46 : Half(1, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%43, %44, %45, %features.5.running_mean, %features.5.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %47 : Half(1, 32, 32, 32, strides=[32768, 1024, 32, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%46) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %48 : Half(1, 32, 16, 16, strides=[8192, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%47) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:719:0
  %49 : Half(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%48, %features.9.weight, %features.9.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %50 : Half(64, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %51 : Half(64, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %52 : Half(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%49, %50, %51, %features.10.running_mean, %features.10.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %53 : Half(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%52) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %54 : Half(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%53, %features.13.weight, %features.13.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %55 : Half(64, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %56 : Half(64, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %57 : Half(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%54, %55, %56, %features.14.running_mean, %features.14.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %58 : Half(1, 64, 16, 16, strides=[16384, 256, 16, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%57) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %59 : Half(1, 64, 8, 8, strides=[4096, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%58) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:719:0
  %60 : Half(1, 128, 8, 8, strides=[8192, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%59, %features.18.weight, %features.18.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %61 : Half(128, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %62 : Half(128, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %63 : Half(1, 128, 8, 8, strides=[8192, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%60, %61, %62, %features.19.running_mean, %features.19.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %64 : Half(1, 128, 8, 8, strides=[8192, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%63) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %65 : Half(1, 128, 8, 8, strides=[8192, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[1, 1, 1, 1], strides=[1, 1]](%64, %features.22.weight, %features.22.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %66 : Half(128, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %67 : Half(128, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %68 : Half(1, 128, 8, 8, strides=[8192, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%65, %66, %67, %features.23.running_mean, %features.23.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %69 : Half(1, 128, 8, 8, strides=[8192, 64, 8, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%68) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %70 : Half(1, 128, 4, 4, strides=[2048, 16, 4, 1], requires_grad=1, device=cuda:0) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%69) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:719:0
  %71 : Half(1, 256, 2, 2, strides=[1024, 4, 2, 1], requires_grad=1, device=cuda:0) = onnx::Conv[dilations=[1, 1], group=1, kernel_shape=[3, 3], pads=[0, 0, 0, 0], strides=[1, 1]](%70, %features.27.weight, %features.27.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/modules/conv.py:490:0
  %72 : Half(256, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %73 : Half(256, strides=[1], device=cpu) = onnx::Constant[value=&lt;Tensor&gt;]()
  %74 : Half(1, 256, 2, 2, strides=[1024, 4, 2, 1], requires_grad=1, device=cuda:0) = onnx::BatchNormalization[epsilon=1.0000000000000001e-05, momentum=0.90000000000000002](%71, %72, %73, %features.28.running_mean, %features.28.running_var) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:2282:0
  %75 : Half(1, 256, 2, 2, strides=[1024, 4, 2, 1], requires_grad=1, device=cuda:0) = onnx::Relu(%74) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1169:0
  %76 : Half(1, 256, 1, 1, strides=[256, 1, 1, 1], requires_grad=1, device=cuda:0) = onnx::MaxPool[kernel_shape=[2, 2], pads=[0, 0, 0, 0], strides=[2, 2]](%75) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:719:0
  %77 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={1}]()
  %78 : Long(requires_grad=0, device=cpu) = onnx::Constant[value={-1}]()
  %79 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[axes=[0]](%77)
  %80 : Long(1, strides=[1], device=cpu) = onnx::Unsqueeze[axes=[0]](%78)
  %81 : Long(2, strides=[1], device=cpu) = onnx::Concat[axis=0](%79, %80)
  %82 : Half(1, 256, strides=[256, 1], requires_grad=1, device=cuda:0) = onnx::Reshape(%76, %81)
  %83 : Half(1, 10, strides=[10, 1], requires_grad=1, device=cuda:0) = onnx::Gemm[alpha=1., beta=1., transB=1](%82, %classifier.0.weight, %classifier.0.bias) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1848:0
  %output : Half(1, 10, strides=[10, 1], requires_grad=1, device=cuda:0) = onnx::Softmax[axis=1](%83) # /home/jordan/.local/lib/python3.8/site-packages/torch/nn/functional.py:1680:0
  return (%output)
</code></pre>
<p>So my question is how can I access these tensors in my pytorch model and force them to gpu? I tried messing with the model’s <code>_apply</code> function as described <a href="https://github.com/ultralytics/yolov5/pull/5110" rel="noopener nofollow ugc">here</a>, but still couldn’t get it to work.</p>

username: ptrblck
cooked: <p>I think the issue you are seeing might be created by <code>torch.jit.trace</code> which is used by the ONNX exporter as the default. Could you try to trace the model without ONNX and see if it’s still failing?<br>
If so, try to use <code>torch.jit.script</code> instead and see if this would be working. If so, use it in the ONNX export as well.</p>

username: Jordan_Chipka
cooked: <p>I tried to use <code>torch.jit.trace</code> and it was able to successfully trace the model.</p>
<pre><code class="lang-auto">traced_module = torch.jit.trace(model, input_tensor) # success
</code></pre>
<p>But I still get the same error when I try to use the traced module in the onnx export.</p>
<pre><code class="lang-auto">    torch.onnx.export(
        traced_module,
        input_tensor,
        onnx_file_path,
        input_names=["input"],
        output_names=["output"],
        export_params=True,
    ) # fail, same as before
</code></pre>

username: ptrblck
cooked: <p>Thanks for the test!<br>
I don’t see any information if <code>jit.script</code> would work in ONNX, but in any case you might want to create an <a href="https://github.com/pytorch/pytorch/issues/">issue</a> so that the code owners could debug the error.</p>

username: Jordan_Chipka
cooked: <p>Thanks for the replies. I tried <code>jit.script</code> in ONNX as well, but still got the same error.<br>
I just opened as issue <a href="https://github.com/pytorch/pytorch/issues/72175" rel="noopener nofollow ugc">here</a>.</p>