[hotfix] add recompile after graph manipulatation (#1621)

2022-09-23 11:00:33 +08:00 · 2022-09-23 11:00:33 +08:00 · 9eae855408
parent d967779a32
commit 9eae855408
1 changed files with 1 additions and 0 deletions
--- a/tests/test_auto_parallel/test_shape_consistency_pass.py
+++ b/tests/test_auto_parallel/test_shape_consistency_pass.py
@ -65,6 +65,7 @@ def check_apply(rank, world_size, port):
    solution = list(ret[0])
    sharding_spec_dict, origin_spec_dict = solution_annotatation_pass(gm, solution, device_mesh)
    shape_consistency_pass(gm)
+    gm.recompile()
    nodes = [node for node in gm.graph.nodes]
    # TODO: wrap the gm to avoid the influence of the user training code
    output = gm(input, sharding_spec_dict, origin_spec_dict)