TorchScript trace in Python while model on GPU gives CUDA error in C++

username: NumesSanguis
cooked: <h1>Python</h1>
<p>With PyTorch (Lightning) I save a model during training to TorchScript with Trace in a Callback (hence the model is on the GPU):</p>
<pre><code class="lang-python">class SimpleModel(pl.LightningModule):
    def __init__(self, **kwargs):
        super().__init__()
        self.layer = torch.nn.Linear(10, 3)

model = SimpleModel()
# in training:
    torch.jit.trace(func=model.eval(), example_inputs=example_inputs)
</code></pre>
<h1>C++</h1>
<p>Now I want to load that trained model in C++ with:</p>
<pre><code class="lang-cpp">torch::jit::script::Module module;
try {
    // Deserialize the ScriptModule from a file using torch::jit::load().
    // torch::Device device(torch::kCPU);
    module = torch::jit::load(argv[1]); //-&gt;to(at::kCPU);  // , map_location=torch::device("cpu")
}
catch (const c10::Error&amp; e) {
    std::cerr &lt;&lt; "error loading the model\n" &lt;&lt; e.what();
    return -1;
}
</code></pre>
<p>which gives me the error:</p>
<pre><code class="lang-auto">error loading the model
Could not run 'aten::empty_strided' with arguments from the 'CUDA' backend. 'aten::empty_strided' is only available for these backends: [CPU, BackendSelect, Autograd, Profiler, Tracer].
Exception raised from reportError at ../aten/src/ATen/core/dispatch/Dispatcher.cpp:313 (most recent call first):
...
</code></pre>
<h2>Traceback</h2>
<details><summary>Full Traceback (CLICK ME)</summary>
<p>
</p>
<pre><code class="lang-auto">$ ./torchscript_simple ../../test.pt
ok
error loading the model
Could not run 'aten::empty_strided' with arguments from the 'CUDA' backend. 'aten::empty_strided' is only available for these backends: [CPU, BackendSelect, Autograd, Profiler, Tracer].
Exception raised from reportError at ../aten/src/ATen/core/dispatch/Dispatcher.cpp:313 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt;) + 0x69 (0x7f516d348eb9 in /project_path/cpp_simple/libtorch/lib/libc10.so)
frame #1: c10::Dispatcher::reportError(c10::DispatchTable const&amp;, c10::DispatchKey) + 0x3dc (0x7f515dc36a5c in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #2: &lt;unknown function&gt; + 0x11a8851 (0x7f515e2cf851 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #3: &lt;unknown function&gt; + 0x11a9e45 (0x7f515e2d0e45 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #4: at::empty_strided(c10::ArrayRef&lt;long&gt;, c10::ArrayRef&lt;long&gt;, c10::TensorOptions const&amp;) + 0x10b (0x7f515e3e5cdb in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #5: &lt;unknown function&gt; + 0x2e4603d (0x7f515ff6d03d in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #6: &lt;unknown function&gt; + 0x11a9e45 (0x7f515e2d0e45 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #7: at::empty_strided(c10::ArrayRef&lt;long&gt;, c10::ArrayRef&lt;long&gt;, c10::TensorOptions const&amp;) + 0x10b (0x7f515e3e5cdb in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #8: at::native::to(at::Tensor const&amp;, c10::Device, c10::ScalarType, bool, bool, c10::optional&lt;c10::MemoryFormat&gt;) + 0x49f (0x7f515e00cecf in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #9: &lt;unknown function&gt; + 0x135a2d1 (0x7f515e4812d1 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #10: &lt;unknown function&gt; + 0x2ee0e01 (0x7f5160007e01 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #11: &lt;unknown function&gt; + 0x13a2eca (0x7f515e4c9eca in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #12: at::Tensor::to(c10::Device, c10::ScalarType, bool, bool, c10::optional&lt;c10::MemoryFormat&gt;) const + 0x146 (0x7f515e541eb6 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #13: &lt;unknown function&gt; + 0x3722871 (0x7f5160849871 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #14: &lt;unknown function&gt; + 0x3723940 (0x7f516084a940 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #15: &lt;unknown function&gt; + 0x3723ef1 (0x7f516084aef1 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #16: torch::jit::readArchiveAndTensors(std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const&amp;, c10::optional&lt;std::function&lt;c10::StrongTypePtr (c10::QualifiedName const&amp;)&gt; &gt;, c10::optional&lt;std::function&lt;c10::intrusive_ptr&lt;c10::ivalue::Object, c10::detail::intrusive_target_default_null_type&lt;c10::ivalue::Object&gt; &gt; (c10::StrongTypePtr, c10::IValue)&gt; &gt;, c10::optional&lt;c10::Device&gt;, caffe2::serialize::PyTorchStreamReader&amp;) + 0x6b2 (0x7f51607eb642 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #17: &lt;unknown function&gt; + 0x36c495d (0x7f51607eb95d in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #18: &lt;unknown function&gt; + 0x36c71ff (0x7f51607ee1ff in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #19: torch::jit::load(std::unique_ptr&lt;caffe2::serialize::ReadAdapterInterface, std::default_delete&lt;caffe2::serialize::ReadAdapterInterface&gt; &gt;, c10::optional&lt;c10::Device&gt;, std::unordered_map&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt;, std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt;, std::hash&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; &gt;, std::equal_to&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; &gt;, std::allocator&lt;std::pair&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const, std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; &gt; &gt; &gt;&amp;) + 0x179 (0x7f51607ee7a9 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #20: torch::jit::load(std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const&amp;, c10::optional&lt;c10::Device&gt;, std::unordered_map&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt;, std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt;, std::hash&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; &gt;, std::equal_to&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; &gt;, std::allocator&lt;std::pair&lt;std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; const, std::__cxx11::basic_string&lt;char, std::char_traits&lt;char&gt;, std::allocator&lt;char&gt; &gt; &gt; &gt; &gt;&amp;) + 0x75 (0x7f51607f1355 in /project_path/cpp_simple/libtorch/lib/libtorch_cpu.so)
frame #21: main + 0xd6 (0x557c5ac9c4df in ./torchscript_simple)
frame #22: __libc_start_main + 0xe7 (0x7f515c5b4b97 in /lib/x86_64-linux-gnu/libc.so.6)
frame #23: _start + 0x2a (0x557c5ac9bb5a in ./torchscript_simple)
</code></pre>

</details>
<h1>Loading model onto a specific device?</h1>
<p>It seems <code>C++</code> does not have an argument like <code>map_location</code> as in Python: <code>torch.jit.load(ts_location, map_location=sample.device)</code></p>
<h1>Question</h1>
<p>How should I load the model on CPU avoiding a CUDA error (without needing to transfer the model to the CPU in Python when saving as TorchScript)?</p>

username: NumesSanguis
cooked: <p>Currently I’m using the workaround where I once again load the TorchScript into Python, but map it onto the cpu, then save it. The new model can successfully be used in C++</p>
<pre><code class="lang-python">import torch

# load TorchScript model on cpu
torchmodel = torch.jit.load("path/to/model.pt", map_location='cpu')

# save, but with no CUDA code
torch.jit.save(torchmodel, "path/to/model_cpu.pt")
</code></pre>
<p>Still I would like to avoid this workaround if possible.</p>

username: fermat97
cooked: <p>Hi, is there any stable solution to this problem? I mean sending torchscript models from GPU to cpu and vice versa, without any saving and loading process.</p>