[SOLVED} Calling symbolic_trace after inserting a call_module node creates problems

nscotto: <p>I am getting errors during a call to <code>symbolic_trace</code> on a <code>GraphModule</code> after having done modifications to it (inserting a <code>call_module</code> node).<br>
This is the exact error::</p>
<pre><code class="lang-auto">/home/SERILOCAL/n.perto/.anaconda3/envs/test/lib/python3.7/site-packages/torch/fx/graph.py:606: UserWarning: Attempted to insert a call_module Node with no underlying reference in the owning GraphModule! Call GraphModule.add_submodule to add the necessary submodule
  warnings.warn("Attempted to insert a call_module Node with "
Traceback (most recent call last):
  File "pytorch_issue.py", line 44, in &lt;module&gt;
    new_mod = fx.symbolic_trace(mod)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/test/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 859, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/test/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 571, in trace
    self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
  File "&lt;eval_with_key_2&gt;", line 3, in forward
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/test/lib/python3.7/site-packages/torch/fx/graph_module.py", line 513, in wrapped_call
    raise e.with_traceback(None)
AttributeError: 'MyOtherModule' object has no attribute 'conv'
</code></pre>
<p>And here is the code:</p>
<pre><code class="lang-auto">import torch
import torch.fx as fx
import torch.nn as nn


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x


class MyOtherModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(24, 32, 3)

    def forward(self, x):
        x = nn.functional.relu(self.conv(x))


def insert_mod_after(mod: fx.GraphModule, from_idx: int, insertme: fx.GraphModule):
    mod.add_submodule('inserted', insertme)
    cutoff_node: fx.Node = list(mod.graph.nodes)[from_idx]
    next_node: fx.Node = list(mod.graph.nodes)[from_idx+1]

    with mod.graph.inserting_after(cutoff_node):
        new_node = mod.graph.call_module('inserted', (cutoff_node,), {})

    next_node.replace_input_with(cutoff_node, new_node)

    mod.delete_all_unused_submodules()
    mod.graph.eliminate_dead_code()
    mod.recompile()
    mod.graph.lint()

    return mod


mod = fx.symbolic_trace(MyModule())
other = fx.symbolic_trace(MyOtherModule())
insert_mod_after(mod, 0, other)
new_mod = fx.symbolic_trace(mod)
</code></pre>
<p>It looks like it could be related to the warning, although I don’t understand why this warning is happening since I am calling <code>add_submodule</code> already… I created <a href="https://discuss.pytorch.org/t/warning-attempted-to-insert-a-call-module-node-with-no-underlying-reference-in-the-owning-graphmodule/125858">another topic</a> for the warning.</p>
<p>Thank you for reading</p>

nscotto: <p>It seems that the problems come from the call to <code>mod.delete_all_unused_submodules()</code>.<br>
I think it deletes submodules  of the <code>Sequential</code> submodule added as they don’t appear in the graph even though the <code>call_module</code> node calling the <code>Sequential</code> submodule use them internally.<br>
Putting it after the second <code>symbolic_trace</code> seems to solve the issue.</p>
<pre><code class="lang-auto">import torch
import torch.fx as fx
import torch.nn as nn


class MyModule(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x


class MyOtherModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(24, 32, 3)

    def forward(self, x):
        x = nn.functional.relu(self.conv(x))


def insert_mod_after(mod: fx.GraphModule, from_idx: int, insertme: fx.GraphModule):
    mod.add_submodule('inserted', insertme)
    cutoff_node: fx.Node = list(mod.graph.nodes)[from_idx]
    next_node: fx.Node = list(mod.graph.nodes)[from_idx+1]

    with mod.graph.inserting_after(cutoff_node):
        new_node = mod.graph.call_module('inserted', (cutoff_node,), {})

    next_node.replace_input_with(cutoff_node, new_node)

    mod.graph.eliminate_dead_code()
    mod.recompile()
    mod.graph.lint()
    mod = fx.symbolic_trace(mod)
    mod.delete_all_unused_submodules()
    return mod


mod = fx.symbolic_trace(MyModule())
other = fx.symbolic_trace(MyOtherModule())
new_mod = insert_mod_after(mod, 0, other)
</code></pre>
<p>I would like these type of ordering to appear in the documentation as it is not clear in which order to call <code>Graph.eliminate_dead_code()</code>, <code>Graph.lint()</code>, <code>GraphModule.delete_all_unused_submodules()</code>, <code>GraphModule.recompile()</code>.</p>