test(AutoParallel): test model add auto_parallel config

Oneflow-Inc · Aug 15, 2022 · 9dbb458 · 9dbb458
1 parent 755d3e6
commit 9dbb458
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/oneflow/test/graph/test_alexnet_auto_parallel.py b/python/oneflow/test/graph/test_alexnet_auto_parallel.py
@@ -100,6 +100,10 @@ class AlexNetEvalGraph(flow.nn.Graph):
         def __init__(self):
             super().__init__()
             self.alexnet = alexnet_module
+            self.config.enable_auto_parallel(True)
+            self.config.enable_auto_parallel_prune_parallel_cast_ops(True)
+            self.config.enable_auto_parallel_mainstem_algo(True)
+            self.config.enable_auto_parallel_sbp_collector(True)
 
         def build(self, image):
             with flow.no_grad():