def nchw2nhwc_pass(
    module: torch.fx.GraphModule,
) -> torch.fx.GraphModule:
    """
    This pass is a kind of hacky way to support some vision models. The reason is due the fact that the frontend is traced based on channel first while AIT needs channel last.
    We need to modify
    1) mean.dim for dim=[-1,-2] changed to [-2,-3]
    2) dim=1 of mean.dim changed to dim=3
    3) concat(inputs, dim=1) need to be dim=3
    """
    modified = False
    for n in module.graph.nodes:
        if n.op == "call_function" and n.target == torch.ops.aten.cat.default:
            if n.args[1] == 1:
                new_args = list(n.args)
                new_args[1] = 3
        elif n.op == "call_function" and n.target == torch.ops.aten.mean.dim:
            if n.args[1] == [-1, -2] or [-2, -1]:
                new_args = list(n.args)
                new_args[1] = [-2, -3]
        else:
            continue
        n.args = tuple(new_args)
        modified = True

        modified_list1 = []
        modified_list2 = []
        modified_list3 = []
        for u in n.users:
            if u.target == torch.ops.aten.sym_size and u.args[1] == 1:
                modified_list1.append(u)
            if u.target == torch.ops.aten.sym_size and u.args[1] == 2:
                modified_list2.append(u)
            if u.target == torch.ops.aten.sym_size and u.args[1] == 3:
                modified_list3.append(u)

        for v in modified_list1:
            new_args = list(v.args)
            new_args[1] = 3
            v.args = tuple(new_args)

        for v in modified_list2:
            new_args = list(v.args)
            new_args[1] = 1
            v.args = tuple(new_args)

        for v in modified_list3:
            new_args = list(v.args)
            new_args[1] = 2
            v.args = tuple(new_args)

    module.graph.eliminate_dead_code()
    module.recompile()
    return PassResult(module, modified)
This pass is a kind of hacky way to support some vision models. The reason is due the fact that the frontend is traced based on channel first while AIT needs channel last.
We need to modify
1) mean.dim for dim=[-1,-2] changed to [-2,-3]
2) dim=1 of mean.dim changed to dim=3
3) concat(inputs, dim=1) need to be dim=3
fx2ait/fx2ait/passes/lower_basic_pass_aten.py