PyBind11 doesn't recognize `torch._C.ScriptModule` as `torch::jit::Module`

Benjamin_Bolte1: <p>Hi! Maybe this isn’t the right place for this question. I am trying to write a PyBind11 Torch extension that takes a ScriptModule as input. I have some function like this:</p>
<pre><code class="lang-auto">torch::jit::script::Module preprocess(const torch::jit::script::Module&amp; mod) {
  torch::jit::script::Module new_mod(mod._ivalue()-&gt;name() + "_txp");
  return new_mod;
}

void add_submodule(pybind11::module_&amp; m) {
  m.def("preprocess",
        &amp;preprocess,
        "Preprocesses the TorchScript module",
        "model"_a);
}
</code></pre>
<p>This compiles fine, but when I call the function from Python I get this error:</p>
<pre><code class="lang-auto">TypeError: func_name(): incompatible function arguments. The following argument types are supported:
    1. (model: torch::jit::Module) -&gt; torch::jit::Module

Invoked with: &lt;torch._C.ScriptModule object at 0x7fc27d657770&gt;
</code></pre>
<p>How do I resolve this?</p>

Benjamin_Bolte1: <p>More context, this works fine:</p>
<pre><code class="lang-auto">torch::Tensor preprocess(const torch::Tensor&amp; mod) { return mod; }

PYBIND11_MODULE(MODULE_NAME, m) { m.def("preprocess", &amp;preprocess); }
</code></pre>
<p>Here’s my CMake file:</p>
<pre><code class="lang-auto">cmake_minimum_required(VERSION 3.12 FATAL_ERROR)

project(txp-cpp-backend LANGUAGES CXX C)

set(LIB_NAME cpp)

set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wformat")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Werror=cpp")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wformat-security")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fstack-protector")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DMODULE_NAME=lib${LIB_NAME}")
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -O3 -Wno-reorder")
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)

find_package(Torch REQUIRED)
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
find_package(PythonLibs REQUIRED)

find_library(TORCH_PYTHON_LIBRARY torch_python
             PATHS "${TORCH_INSTALL_PREFIX}/lib")
message(STATUS "TORCH_PYTHON_LIBRARY: ${TORCH_PYTHON_LIBRARY}")

add_library(${LIB_NAME} SHARED txp.cpp)

target_compile_features(${LIB_NAME} PRIVATE cxx_std_14)
set_target_properties(${LIB_NAME} PROPERTIES OUTPUT_NAME ${LIB_NAME})
set_target_properties(${LIB_NAME} PROPERTIES POSITION_INDEPENDENT_CODE ON)
set_target_properties(${LIB_NAME}
                      PROPERTIES LINK_FLAGS "-Wl,-rpath,${CMAKE_INSTALL_RPATH}")
target_compile_options(
  ${LIB_NAME} PUBLIC -Wall -Wno-sign-compare -Wno-unused-function
                     -Wno-unknown-pragmas)
target_include_directories(
  ${LIB_NAME}
  PUBLIC ${TORCH_INCLUDE_DIRS}
  PUBLIC ${PYTHON_INCLUDE_DIRS})
target_link_libraries(
  ${LIB_NAME}
  PUBLIC ${PYTHON_LIBRARIES}
  PUBLIC ${TORCH_LIBRARIES}
  PUBLIC ${TORCH_PYTHON_LIBRARY})
</code></pre>
<p>And this guy:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; import torch; torch.__version__
'1.11.0.dev20211121'
</code></pre>
<p>When I try to add <code>#include &lt;torch/csrc/jit/python/script_init.h&gt;</code> I get:</p>
<pre><code class="lang-auto">fatal error: torch/csrc/generic/Storage.h: No such file or directory
</code></pre>
<p>which is indeed missing from the Torch install path</p>

Michael_Suo: <p>What does the Python code look like? What I suspect is happening is: torch.jit.ScriptModule<code>is a Python wrapper; the actual pybinded type is actually stored in the attribute</code>_c` on it.</p>

Benjamin_Bolte1: <p>Something like this:</p>
<pre><code class="lang-auto">script_model = torch.jit.load(cfg.model_file, map_location="cpu")
preprocessed_model = cpplib.preprocess(script_model._c)
</code></pre>
<p>Yep, I am using the <code>._c</code> guy rather than the original module</p>

Benjamin_Bolte1: <p>Ok, turns out I was missing some important <code>-DPYBIND11_</code> directives. I’ve created a minimal example repo here for posterity’s sake: <a href="https://github.com/codekansas/torchscript-example" class="inline-onebox" rel="noopener nofollow ugc">GitHub - codekansas/torchscript-example: Example CMake project for TorchScript</a></p>