def replace_batch_norm(module: torch.fx.GraphModule) -> torch.fx.GraphModule:
    """
    Current exir.capture enable_aot and it captures bwd needed nodes in output P619801318
    This pass removes those unused node and replace with classic aten.batch_norm
    """
    batch_node_list = []
    for n in module.graph.nodes:
        if n.target == torch.ops.aten._native_batch_norm_legit_functional.default:
            batch_node_list.append(n)
        if n.target == "output":
            output_node = n

    if len(batch_node_list) > 0:
        modified = True
    else:
        modified = False
    for n in batch_node_list:
        new_op = torch.ops.aten.batch_norm
        new_args = list(n.args)
        new_args.append(False)
        new_args = tuple(new_args)
        user_list = [x for x in n.users]
        user_list_copy_node = []
        user_list_copy_node.append(next(iter(user_list[1].users)))
        user_list_copy_node.append(next(iter(user_list[2].users)))
        getitem_node = user_list[0]
        with module.graph.inserting_after(getitem_node):
            new_node = module.graph.create_node(
                "call_function",
                new_op,
                args=new_args,
                kwargs=n.kwargs,
            )
            getitem_node.replace_all_uses_with(new_node)

        output_args = output_node.args[0]
        new_output_args = [x for x in output_args if x not in user_list_copy_node]
        output_node.args = (new_output_args,)
        module.graph.eliminate_dead_code()
        module.recompile()
    return PassResult(module, modified)
Current exir.capture enable_aot and it captures bwd needed nodes in output P619801318
This pass removes those unused node and replace with classic aten.batch_norm
fx2ait/fx2ait/passes/lower_basic_pass_aten.py