Unable to JIT trace Tacotron 2 (from `torchaudio` library)

rahulbhalley: <p>Hi everyone!</p>
<p><em><strong>I am unable to trace the Tacotron 2 model given in <code>torchaudio</code> library.</strong></em></p>
<p><strong>Here’s my approach</strong></p>
<pre><code class="lang-python">import torch
from torch import nn
import torchaudio

device = "cuda" if torch.cuda.is_available() else "cpu"

bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_PHONE_LJSPEECH
processor = bundle.get_text_processor()
tacotron2 = bundle.get_tacotron2().to(device)

text = "Hello world! Text to speech!"
processed, lengths = processor(text)
</code></pre>
<p><strong>Let’s define a model wrapper for tracing</strong></p>
<pre><code class="lang-python">class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = tacotron2
    def forward(self, processed, lengths):
        spec, _, _ = self.model.infer(processed, lengths)
        return spec

model = Model()
spec = model(processed, lengths)
print(spec.shape) # torch.Size([1, 80, 193])
</code></pre>
<h3>
<a name="lets-try-to-trace-it-1" class="anchor" href="#lets-try-to-trace-it-1"></a>Let’s try to trace it</h3>
<pre><code class="lang-python">traced_model = torch.jit.trace(model, (processed, lengths))
</code></pre>
<p><strong>Now I get the error</strong></p>
<pre><code class="lang-auto">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[29], line 1
----&gt; 1 traced_model = torch.jit.trace(model, (processed, lengths))

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_trace.py:794, in trace(func, example_inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_kwarg_inputs, _store_inputs)
    792         else:
    793             raise RuntimeError("example_kwarg_inputs should be a dict")
--&gt; 794     return trace_module(
    795         func,
    796         {"forward": example_inputs},
    797         None,
    798         check_trace,
    799         wrap_check_inputs(check_inputs),
    800         check_tolerance,
    801         strict,
    802         _force_outplace,
    803         _module_class,
    804         example_inputs_is_kwarg=isinstance(example_kwarg_inputs, dict),
    805         _store_inputs=_store_inputs
    806     )
    807 if (
    808     hasattr(func, "__self__")
    809     and isinstance(func.__self__, torch.nn.Module)
    810     and func.__name__ == "forward"
    811 ):
    812     if example_inputs is None:

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_trace.py:1023, in trace_module(mod, inputs, optimize, check_trace, check_inputs, check_tolerance, strict, _force_outplace, _module_class, _compilation_unit, example_inputs_is_kwarg, _store_inputs)
   1020 torch.jit._trace._trace_module_map = trace_module_map
   1021 register_submods(mod, "__module")
-&gt; 1023 module = make_module(mod, _module_class, _compilation_unit)
   1025 for method_name, example_inputs in inputs.items():
   1026     if method_name == "forward":
   1027         # "forward" is a special case because we need to trace
   1028         # `Module.__call__`, which sets up some extra tracing, but uses
   1029         # argument names of the real `Module.forward` method.

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_trace.py:604, in make_module(mod, _module_class, _compilation_unit)
    602 if _module_class is None:
    603     _module_class = TopLevelTracedModule
--&gt; 604 return _module_class(mod, _compilation_unit=_compilation_unit)

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_trace.py:1169, in TracedModule.__init__(self, orig, id_set, _compilation_unit)
   1167     if submodule is None:
   1168         continue
-&gt; 1169     tmp_module._modules[name] = make_module(
   1170         submodule, TracedModule, _compilation_unit=None
   1171     )
   1173 script_module = torch.jit._recursive.create_script_module(
   1174     tmp_module, lambda module: (), share_types=False, is_tracing=True
   1175 )
   1177 self.__dict__["_name"] = type(orig).__name__

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_trace.py:595, in make_module(mod, _module_class, _compilation_unit)
    592 elif torch._jit_internal.module_has_exports(mod):
    594     infer_methods_stubs_fn = torch.jit._recursive.make_stubs_from_exported_methods
--&gt; 595     return torch.jit._recursive.create_script_module(
    596         mod,
    597         infer_methods_stubs_fn,
    598         share_types=False,
    599         is_tracing=True
    600     )
    601 else:
    602     if _module_class is None:

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_recursive.py:480, in create_script_module(nn_module, stubs_fn, share_types, is_tracing)
    478 if not is_tracing:
    479     AttributeTypeIsSupportedChecker().check(nn_module)
--&gt; 480 return create_script_module_impl(nn_module, concrete_type, stubs_fn)

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_recursive.py:546, in create_script_module_impl(nn_module, concrete_type, stubs_fn)
    544 # Compile methods if necessary
    545 if concrete_type not in concrete_type_store.methods_compiled:
--&gt; 546     create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    547     # Create hooks after methods to ensure no name collisions between hooks and methods.
    548     # If done before, hooks can overshadow methods that aren't exported.
    549     create_hooks_from_stubs(concrete_type, hook_stubs, pre_hook_stubs)

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_recursive.py:397, in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    394 property_defs = [p.def_ for p in property_stubs]
    395 property_rcbs = [p.resolution_callback for p in property_stubs]
--&gt; 397 concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_recursive.py:898, in compile_unbound_method(concrete_type, fn)
    894 stub = make_stub(fn, fn.__name__)
    895 with torch._jit_internal._disable_emit_hooks():
    896     # We don't want to call the hooks here since the graph that is calling
    897     # this function is not yet complete
--&gt; 898     create_methods_and_properties_from_stubs(concrete_type, (stub,), ())
    899 return stub

File /opt/homebrew/lib/python3.9/site-packages/torch/jit/_recursive.py:397, in create_methods_and_properties_from_stubs(concrete_type, method_stubs, property_stubs)
    394 property_defs = [p.def_ for p in property_stubs]
    395 property_rcbs = [p.resolution_callback for p in property_stubs]
--&gt; 397 concrete_type._create_methods_and_properties(property_defs, property_rcbs, method_defs, method_rcbs, method_defaults)

RuntimeError: Couldn't find method: 'forward__0' on class: '__torch__.torch.nn.modules.rnn.LSTM (of Python compilation unit at: 0x6000007dcb58)'
</code></pre>
<p><strong>Main error line is</strong></p>
<pre><code class="lang-auto">RuntimeError: Couldn't find method: 'forward__0' on class: '__torch__.torch.nn.modules.rnn.LSTM (of Python compilation unit at: 0x6000007dcb58)'
</code></pre>
<p>Could any please help me solve this?</p>
<p>Thanks,<br>
Rahul Bhalley</p>