Help with custom c++ class in a jit model

username: dma_science
cooked: <p>Hi all,<br>
I am having some difficulties getting my custom c++ class to work correctly with torch::jit::load despite being able to use it in python.</p>
<p>Essentially I have a custom class that follows this structure in the cpp file</p>
<pre><code class="lang-auto">TORCH_LIBRARY(mycustomop, m) {
  m.class_&lt;CustomComputer&gt;("CustomComputer")
      .def(torch::init&lt; Tensor, Tensor, Tensor, Tensor &gt;())
      .def_pickle(
          // __getstate__
          [](const c10::intrusive_ptr&lt;CustomComputer&gt;&amp; self) -&gt; std::vector&lt;Tensor&gt; {
            std::vector&lt;Tensor&gt; state;
            state.push_back(self-&gt; params.a);
            state.push_back(self-&gt;params.b);
            state.push_back(self-&gt;params.c);
            state.push_back(self-&gt;params.d);
            return state;
          },
          // __setstate__
          [](std::vector&lt;Tensor&gt; state) -&gt; c10::intrusive_ptr&lt;CustomComputer&gt; {
            return c10::make_intrusive&lt;CustomComputer&gt;(
                state[0],
                state[1],
                state[2],
                state[3]);
          });
</code></pre>
<p>however when I attempt to jit load a model that uses this class I run into the follow exception</p>
<pre><code class="lang-auto">Exception:
Unknown type name '__torch__.torch.classes.mycustomop.CustomComputer':
</code></pre>
<p>I have no problem using this model on the python side if I do something like</p>
<pre><code class="lang-auto">import torch
import mycustomop
model = torch.jit.load('my_compiled_model')
</code></pre>
<p>Can someone help with how to properly register this custom class on the c++ side.</p>

username: dbrenes
cooked: <p>Hello <a class="mention" href="/u/dma_science">@dma_science</a></p>
<p>I was able to reproduce your error on my system:<br>
<code>Unknown type name '__torch__.torch.classes.mycustomop.CustomComputer':</code></p>
<p>I managed to fix it and sucessfully load the model by linking the <code>CustomComputer</code> class during the build of the C++ binaries. In my case, when using CMake, adding the <code>.cpp</code> file with the <code>CustomComputer</code> implementation to the executables allows me to correctly load the model.</p>
<p>Maybe you can share exactly how you’re building your binaries, so we can take a closer look at where’s the issue.</p>
<hr>
<p><em>Machine Learning Engineer at <a href="https://www.ridgerun.ai/" rel="noopener nofollow ugc">RidgeRun.ai</a></em><br>
<em>Contact us: <a href="mailto:support@ridgerun.ai">support@ridgerun.ai</a></em></p>

username: dma_science
cooked: <p>Thank you for the response! My custom class is being used as part of a larger project so the cmakelists.txt is rather long and complicated, i.e., it likely is a linking problem. Is it possible for you to show the solution for your minimal example? That will probably allow me to debug on my end. Thank you!</p>

username: dbrenes
cooked: <p>Sure, no problem. This is my <code>CMakeLists.txt</code>, where <code>computer.cpp</code> contains the implementation of the <code>CustomComputer</code> class as well as the <code>TORCH_LIBRARY</code> definition, which is the same as the one you shared earlier. <code>main.cpp</code> the is where model is loaded with <code>torch::jit::load</code>.</p>
<pre><code class="lang-auto">cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(MyProject)

set(CMAKE_PREFIX_PATH "${CMAKE_CURRENT_SOURCE_DIR}/libtorch")
find_package(Torch REQUIRED)

add_executable(main main.cpp computer.cpp)
target_link_libraries(main "${TORCH_LIBRARIES}")
set_property(TARGET main PROPERTY CXX_STANDARD 17)
</code></pre>
<p>I managed to reproduce the error by removing <code>computer.cpp</code> from the <code>add_executable</code> line, which led me to believe that it stems from linking issues. Hope this helps.</p>