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

hpcaitech · Sep 23, 2022 · 9eae855 · 9eae855
1 parent d967779
commit 9eae855
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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)