def METHOD_NAME(pattern):
    """ Given a node pattern, extract the corresponding modules
    e.g. input: (relu_node, (bn_node, conv_node))
         output: (relu_module, (bn_module, conv_module))
    """
    if isinstance(pattern, (tuple, list)):
        n, *args = pattern
        modules: List[torch.nn.Module] = []
        modules.append(METHOD_NAME(n))
        for a in args:
            modules.append(METHOD_NAME(a))
        return tuple(modules)
    else:
        n = pattern
        if n.op == "call_module":
            return named_modules[n.target]
        elif n.op == "call_function" and n.target == torch.nn.functional.relu:
            relu = torch.nn.ReLU()
            relu.training = root_module.training
            return relu
        elif n.op == "call_function" or n.op == "call_method":
            return n.target
        else:
            return MatchAllNode