diff --git a/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py b/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py
index b29cd265bd64c90009a4d325b190bdc8e9fbc5a8..0280dfcf67b1df041ba964a0e429ed7b89ec729c 100644
--- a/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py
+++ b/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py
@@ -24,6 +24,7 @@ import unittest
 import os
 import copy
 import numpy as np
+from paddle.static.amp import decorate
 
 paddle.enable_static()
 
@@ -138,7 +139,7 @@ def train(net_type, use_cuda, save_dirname, is_local):
 
         amp_lists = fluid.contrib.mixed_precision.AutoMixedPrecisionLists(
             custom_black_varnames={"loss", "conv2d_0.w_0"})
-        mp_optimizer = fluid.contrib.mixed_precision.decorate(
+        mp_optimizer = decorate(
             optimizer=optimizer,
             amp_lists=amp_lists,
             init_loss_scaling=8.0,
@@ -442,7 +443,7 @@ class TestAmpWithNonIterableDataLoader(unittest.TestCase):
                 optimizer = fluid.optimizer.Lamb(learning_rate=0.001)
                 amp_lists = fluid.contrib.mixed_precision.AutoMixedPrecisionLists(
                     custom_black_varnames={"loss", "conv2d_0.w_0"})
-                mp_optimizer = fluid.contrib.mixed_precision.decorate(
+                mp_optimizer = decorate(
                     optimizer=optimizer,
                     amp_lists=amp_lists,
                     init_loss_scaling=8.0,
diff --git a/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py b/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py
index 64ef2e26bbdb94ccb8e3835c1432c9b658b69c9d..83b920642b847e86f2323e652dceadc70200515c 100644
--- a/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py
+++ b/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py
@@ -19,8 +19,8 @@ import paddle.fluid as fluid
 import contextlib
 import unittest
 import numpy as np
-from paddle.fluid.contrib.mixed_precision.fp16_utils import cast_model_to_fp16
-from paddle.fluid.contrib.mixed_precision.fp16_utils import cast_parameters_to_fp16
+from paddle.static.amp import cast_model_to_fp16
+from paddle.static.amp import cast_parameters_to_fp16
 
 paddle.enable_static()
 
diff --git a/python/paddle/static/__init__.py b/python/paddle/static/__init__.py
index 6778149e2bf0f6fb343e5b0f3624ee58467218c0..9c911e722dbc76629443354fdd432aefc5cb96c6 100644
--- a/python/paddle/static/__init__.py
+++ b/python/paddle/static/__init__.py
@@ -24,6 +24,7 @@ __all__ = [
 ]
 
 from . import nn
+from . import amp
 from .io import save_inference_model  #DEFINE_ALIAS
 from .io import load_inference_model  #DEFINE_ALIAS
 from ..fluid import Scope  #DEFINE_ALIAS
diff --git a/python/setup.py.in b/python/setup.py.in
index 34faff6bea524e7ccd59322881b14e574606620f..63a8ca895614201859a58e0d7f1b3a239fce91b7 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -210,6 +210,7 @@ packages=['paddle',
           'paddle.metric',
           'paddle.static',
           'paddle.static.nn',
+          'paddle.static.amp',
           'paddle.tensor',
           'paddle.onnx',
           ]