diff --git a/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_c_setup.py b/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_c_setup.py
index be096e1c31c44b2c03839d31c0f43814bccbcb48..34778e0d5714d42ca93c8dc91d16ee2ab5973830 100644
--- a/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_c_setup.py
+++ b/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_c_setup.py
@@ -13,9 +13,9 @@
 # limitations under the License.
 
 import os
-from distutils.core import Extension, setup
 from distutils.sysconfig import get_python_lib
 
+from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
 from paddle.fluid import core
diff --git a/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_setup.py b/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_setup.py
index 5ec8345812b596b5369662103713b001d2492c83..603e3666f09e9556b9dd48999afa097e17e2db57 100644
--- a/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_setup.py
+++ b/python/paddle/fluid/tests/custom_kernel/custom_kernel_dot_setup.py
@@ -14,8 +14,8 @@
 
 import os
 import site
-from distutils.core import Extension, setup
 
+from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
 from paddle.fluid import core