Hello,
In the example for the replace pattern of torch.fx, function or replaced (torch.add by torch.mul).
This is very clear, however it is not clear to me if it is possible to replace modules as well, and if so, how to do it.
The following example failed for me with the error:
Traceback (most recent call last):
  File "test.py", line 43, in <module>
    torch.fx.subgraph_rewriter.replace_pattern(gm, pattern, replacement)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/subgraph_rewriter.py", line 201, in replace_pattern
    pattern_graph = symbolic_trace(pattern).graph
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 606, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 355, in trace
    self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
  File "test.py", line 33, in pattern
    val1 = F.relu(lin1(x))
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 344, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 221, in call_module
    module_qualified_name = self.path_of_module(m)
  File "/home/SERILOCAL/n.perto/.anaconda3/envs/automl/lib/python3.7/site-packages/torch/fx/symbolic_trace.py", line 191, in path_of_module
    raise NameError('module is not installed as a submodule')
NameError: module is not installed as a submodule

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.fx


class TwoLayerSubNet(nn.Module):
    def __init__(self, D_in=100, H=50, D_out=10):
        super(TwoLayerSubNet, self).__init__()
        self.linear1 = nn.Linear(D_in, H)
        self.linear2 = nn.Linear(H, D_out)

    def forward(self, x):
            h_relu = F.relu(self.linear1(x))
            y_pred = self.linear2(h_relu)
            return y_pred


class TwoLayerNet(nn.Module):
    def __init__(self, D_in=100, H=50, D_int=20, D_out=10):
        super(TwoLayerNet, self).__init__()
        self.sub1 = TwoLayerSubNet(D_in, H, D_int)
        self.sub2 = TwoLayerSubNet(D_int, H, D_out)

    def forward(self, x):
            h_sub1 = self.sub1(x)
            y_pred = self.sub2(x)
            return y_pred

def pattern(x):
    lin1 = torch.nn.Linear(30, 20)
    lin2 = torch.nn.Linear(20, 10)
    val1 = F.relu(lin1(x))
    return F.relu(lin2(val1))


def replacement(x):
    lin1 = torch.nn.Linear(30, 10)
    return F.relu(lin1(x))

m = TwoLayerNet()
gm = torch.fx.symbolic_trace(m)
torch.fx.subgraph_rewriter.replace_pattern(gm, pattern, replacement)

Am I missing something?
Thanks