Loading TorchScript using LibTorch exported using torch.jit.trace_module.()

Stefan1: <p>Dear all,</p>
<p>I’m no C++ expert so forgive me if I ask something simple. But I cannot seem to figure this out.</p>
<p>I have a module for which I wish to export several methods to TorchScript using Python and import them in C++. For this I read from the PyTorch documentation: <a href="https://pytorch.org/docs/stable/generated/torch.jit.trace_module.html#torch-jit-trace-modulel" rel="noopener nofollow ugc">torch.jit.trace_module</a> and <a href="https://pytorch.org/docs/stable/generated/torch.jit.trace.html#torch-jit-trace" rel="noopener nofollow ugc">torch.jit.trace</a>.</p>
<p>When exporting things from Python using <code>torch.jit.trace</code>, I follow the same approach which always works perfectly. This time I really cannot use methods named <code>forward</code> so I have to go though <code>torch.jit.trace_module</code>.</p>
<p>I roughly use this Python code:</p>
<pre><code class="lang-auto">my_module = MyModule()

trace_inputs = {
	'method1' : self.create_dummy_input_method1()
	'method2' : self.create_dummy_input_method2()
	}

traced_torch_script = torch.jit.trace_module(my_module , trace_inputs)
traced_torch_script.save(export_path)
</code></pre>
<p>Exporting works <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"> .</p>
<p>Then, I try to load this from C++ using LibTorch:</p>
<pre><code class="lang-auto">torch::jit::script::Module myModule = torch::jit::load(import_path);
</code></pre>
<p>This also works (at least, I get no error). Then, I try to call a method on the module:</p>
<pre><code class="lang-auto">std::vector&lt;torch::jit::IValue&gt; inputs;
inputs = ...
torch::Tensor output_tensor = module.forward(inputs).toTensor();
</code></pre>
<p>This fails, which seems logical as I did not export a forward method. (this does work if I exported the module using <code>jit::trace</code> instead of <code>jit::trace_module</code>).<br>
However, the <a href="https://pytorch.org/docs/stable/generated/torch.jit.trace_module.html#torch-jit-trace-modulel" rel="noopener nofollow ugc">torch.jit.trace_module docs</a> state:</p>
<blockquote>
<p><strong>TORCH.JIT.TRACE_MODULE<br>
…<br>
Returns</strong><br>
A <a href="https://pytorch.org/docs/stable/generated/torch.jit.ScriptModule.html#torch.jit.ScriptModule" rel="noopener nofollow ugc"><code>ScriptModule</code></a> object with a single <code>forward</code> method containing the traced code. When <code>func</code> is a <code>torch.nn.Module</code> , the returned <a href="https://pytorch.org/docs/stable/generated/torch.jit.ScriptModule.html#torch.jit.ScriptModule" rel="noopener nofollow ugc"><code>ScriptModule</code></a> will have the same set of sub-modules and parameters as <code>func</code> .</p>
</blockquote>
<p>Is this a copy past error in the docs <a href="https://pytorch.org/docs/stable/generated/torch.jit.trace.html" rel="noopener nofollow ugc">from the <code>jit::trace</code> docs</a>? It seems not to be copy past, so no error. Then, I’m affraid I don’t understand it.</p>
<p>I also tried changing <code>forward</code> to <code>method1</code>:</p>
<pre><code class="lang-auto">torch::Tensor output_tensor = module.method1(inputs).toTensor();
</code></pre>
<p>Unfortunately, then I get compile error saying module does not have a <code>method1</code> (which makes sense to me from a C++ point of view).</p>
<p>How can I call <code>MyModule</code>’s methods <code>method1</code> and <code>method2</code> in C++ using LibTorch?</p>
<p>Thank you!</p>

fatima1: <p>It might be too late to answer this question, but I just wanted to share the solution to this issue:</p>
<p>When we use <code>trace_module()</code> to trace methods other than <code>forward</code>, and we are loading and using that in C++, we should first ‘get’ that specific method, and then we can call it directly. For example:</p>
<pre><code class="lang-auto">auto get_diffusion_noise = model.get_method("get_diffusion_noise");
auto out = get_diffusion_noise(inputs).toTensor();
</code></pre>