How to insert `torch.nn.Linear` in the existing graph

yenming: <p>Hi, there</p>
<p>I am new to <code>torch.fx</code> but I like it! I met a problem when I attempted to insert an operator into an existing GraphModule.</p>
<p>If I insert a <code>torch.relu</code> at the end of the graph, the following program runs successfully. But I fail to insert a <code>torch.nn.Linear</code> into the graph.</p>
<pre><code class="lang-auto">from torch.fx import symbolic_trace
import torch.fx as fx
import torch.nn as nn

import torch.fx

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(100, 50)

    def forward(self, x):
        x = self.fc1(x)
        return x

toy_net = Net()
toy_net_gm = torch.fx.symbolic_trace(toy_net)
print(toy_net_gm.graph)
print(toy_net_gm.code)

def manipulate_graph(graph_module, new_function, module_type):
    for last_node in graph_module.graph.nodes:
        if last_node.op == "output":
            new_func_node = graph_module.graph.call_function(new_function, args=(last_node.prev,))
            with graph_module.graph.inserting_before(last_node):
                last_node.replace_all_uses_with(new_func_node)
            break
    else:
        raise Exception("No found output node")
    graph_module.graph.erase_node(last_node)
    graph_module.graph.output(new_func_node, type_expr=getattr(new_func_node, 'type', None))
    graph_module.graph.lint()
    return fx.GraphModule(module_type, graph_module.graph)

new_gm = manipulate_graph(toy_net_gm, torch.relu, toy_net)
new_gm.recompile()
print(new_gm.graph)
print(new_gm.code)
</code></pre>
<p>I tried to insert a <code>torch.nn.Linear(50,10)</code> into the graph but got an exception:</p>
<pre><code class="lang-auto">  File "E:/pyworkspace/pytorch/add_linear.py", line 58, in &lt;module&gt;
    new_gm = manipulate_graph(toy_net_gm, torch.nn.Linear(50,10), toy_net)
  File "E:/pyworkspace/pytorch/add_linear.py", line 47, in manipulate_graph
    new_func_node = graph_module.graph.call_function(new_function, args=(last_node.prev,))
  File "D:\Softwares\anaconda3\envs\library\lib\site-packages\torch\fx\graph.py", line 728, in call_function
    return self.create_node('call_function', the_function, args, kwargs, type_expr=type_expr)
  File "D:\Softwares\anaconda3\envs\library\lib\site-packages\torch\fx\graph.py", line 428, in create_node
    candidate = name if name is not None else self._target_to_str(target)
  File "D:\Softwares\anaconda3\envs\library\lib\site-packages\torch\fx\graph.py", line 784, in _target_to_str
    op = target.__name__
  File "D:\Softwares\anaconda3\envs\library\lib\site-packages\torch\nn\modules\module.py", line 1185, in __getattr__
    raise AttributeError("'{}' object has no attribute '{}'".format(
AttributeError: 'Linear' object has no attribute '__name__'. Did you mean: '__ne__'?
</code></pre>
<p>The question is how to correctly insert a <code>torch.nn.XXX</code> operator in the graph ?(instead of the function like <code>torch.relu</code>)</p>

yenming: <p>Solved!</p>
<ol>
<li>Firstly, use <code>toy_net_gm.add_module("fc1",torch.nn.Linear(50,10))</code> to define a module in <code>toy_net_gm </code>
</li>
<li>Then use <code>graph_module.graph.call_module("fc1", args=(last_node.prev,))</code> to insert a Linear operator</li>
</ol>