diff --git a/python/paddle/distributed/auto_parallel/helper.py b/python/paddle/distributed/auto_parallel/helper.py
index 6bc177efc9de99f0afa9012ed1a5802786a3b660..4a3a1ab5e15cc08c85de22030f956a75ef584e84 100644
--- a/python/paddle/distributed/auto_parallel/helper.py
+++ b/python/paddle/distributed/auto_parallel/helper.py
@@ -341,6 +341,7 @@ class ProgramHelper(object):
         try:
             return self.proxy_layer.startup_program
         except Exception as err:
+            self._logger.warning("`lazy init` failed.")
             if isinstance(err, AssertionError):
                 return self.concrete_program.startup_program
             raise err
diff --git a/python/paddle/distributed/auto_parallel/parallelizer_v2.py b/python/paddle/distributed/auto_parallel/parallelizer_v2.py
index b83a19b512ef8be8c6fc13aede8056676de4483d..a1dd58fef71316507098c2e93aa1f4e403fab6a9 100644
--- a/python/paddle/distributed/auto_parallel/parallelizer_v2.py
+++ b/python/paddle/distributed/auto_parallel/parallelizer_v2.py
@@ -105,9 +105,13 @@ class Parallelizer:
                 format(time.time() - time0, self._mode))
         else:
             # Apply pre optimization passes
-            # self._apply_pre_optimization(serial_main_program,
-            #                              serial_startup_program, None, None,
-            #                              None)
+            time0 = time.time()
+            self._apply_pre_optimization(serial_main_program,
+                                         serial_startup_program, None, None,
+                                         None)
+            self._logger.info(
+                "within parallel apply_pre_optimization time: {}, mode {}".
+                format(time.time() - time0, self._mode))
             # Do logical partition
             time0 = time.time()
             partitioner = Partitioner(self._dist_context, rank)
@@ -176,7 +180,7 @@ class Parallelizer:
         # apply amp pass
         # FIXME we disenable amp for eval since it has a little bug with
         # eval program and which will be fixed in future
-        if self._mode == 'train' and self._strategy.amp.enable:
+        if self._strategy.amp.enable:
             config = copy.deepcopy(self._strategy.amp.to_dict())
             config["dist_context"] = self._dist_context
             config["params_grads"] = params_grads
diff --git a/python/paddle/fluid/lazy_init.py b/python/paddle/fluid/lazy_init.py
index c44ed7f118943446660502ed78fb2b22df2d35a2..8fc175ae3efebf31948f7dd84c2fb6dbc9dbed99 100644
--- a/python/paddle/fluid/lazy_init.py
+++ b/python/paddle/fluid/lazy_init.py
@@ -36,7 +36,7 @@ class LazyInitHelper(object):
         """
         if self._state:
             return
-        assert framework.in_dygraph_mode(
+        assert framework._non_static_mode(
         ), "LazyInit.enable() is only available in dygraph mode."
         self._state = True