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 39d47d6f4483b99f29190084ddc6976875065730..a29852d722ce72ee91074b1c362b7ce95494e482 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,18 +13,19 @@
 # limitations under the License.
 
 import os
-from paddle.fluid import core
 from distutils.sysconfig import get_python_lib
-from distutils.core import setup, Extension
+
+from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
+from paddle.fluid import core
+
 
 # refer: https://note.qidong.name/2018/03/setup-warning-strict-prototypes
 # Avoid a gcc warning below:
 # cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid
 # for C/ObjC but not for C++
 class BuildExt(build_ext):
-
     def build_extensions(self):
         if '-Wstrict-prototypes' in self.compiler.compiler_so:
             self.compiler.compiler_so.remove('-Wstrict-prototypes')
@@ -48,8 +49,9 @@ paddle_custom_kernel_include = [
     os.path.join(site_packages_path, 'paddle', 'include'),
 ]
 # include path third_party
-compile_third_party_path = os.path.join(os.environ['PADDLE_BINARY_DIR'],
-                                        'third_party')
+compile_third_party_path = os.path.join(
+    os.environ['PADDLE_BINARY_DIR'], 'third_party'
+)
 paddle_custom_kernel_include += [
     os.path.join(compile_third_party_path, 'install/gflags/include'),  # gflags
     os.path.join(compile_third_party_path, 'install/glog/include'),  # glog
@@ -69,10 +71,13 @@ custom_kernel_dot_module = Extension(
     include_dirs=paddle_custom_kernel_include,
     library_dirs=paddle_custom_kernel_library_dir,
     libraries=libs,
-    extra_compile_args=paddle_extra_compile_args)
+    extra_compile_args=paddle_extra_compile_args,
+)
 
-setup(name='custom_kernel_dot_c',
-      version='1.0',
-      description='custom kernel fot compiling',
-      cmdclass={'build_ext': BuildExt},
-      ext_modules=[custom_kernel_dot_module])
+setup(
+    name='custom_kernel_dot_c',
+    version='1.0',
+    description='custom kernel fot compiling',
+    cmdclass={'build_ext': BuildExt},
+    ext_modules=[custom_kernel_dot_module],
+)
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 ba116526d3ff4ba8bd7a4148020758fad4d0710b..7fd37db98044af0bff2960d177927a3451ce12c7 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,18 +14,17 @@
 
 import os
 import site
-from paddle.fluid import core
-from distutils.sysconfig import get_python_lib
-from distutils.core import setup, Extension
+from setuptools import Extension, setup
 from setuptools.command.build_ext import build_ext
 
+from paddle.fluid import core
+
 
 # refer: https://note.qidong.name/2018/03/setup-warning-strict-prototypes
 # Avoid a gcc warning below:
 # cc1plus: warning: command line option ‘-Wstrict-prototypes’ is valid
 # for C/ObjC but not for C++
 class BuildExt(build_ext):
-
     def build_extensions(self):
         if '-Wstrict-prototypes' in self.compiler.compiler_so:
             self.compiler.compiler_so.remove('-Wstrict-prototypes')
@@ -46,12 +45,15 @@ if core.is_compiled_with_npu():
 # include path
 site_packages_path = site.getsitepackages()
 paddle_custom_kernel_include = list(
-    map(lambda path: os.path.join(path, 'paddle', 'include'),
-        site_packages_path))
+    map(
+        lambda path: os.path.join(path, 'paddle', 'include'), site_packages_path
+    )
+)
 
 # include path third_party
-compile_third_party_path = os.path.join(os.environ['PADDLE_BINARY_DIR'],
-                                        'third_party')
+compile_third_party_path = os.path.join(
+    os.environ['PADDLE_BINARY_DIR'], 'third_party'
+)
 paddle_custom_kernel_include += [
     os.path.join(compile_third_party_path, 'install/gflags/include'),  # gflags
     os.path.join(compile_third_party_path, 'install/glog/include'),  # glog
@@ -59,7 +61,8 @@ paddle_custom_kernel_include += [
 
 # libs path
 paddle_custom_kernel_library_dir = list(
-    map(lambda path: os.path.join(path, 'paddle', 'fluid'), site_packages_path))
+    map(lambda path: os.path.join(path, 'paddle', 'fluid'), site_packages_path)
+)
 
 # libs
 libs = [':libpaddle.so']
@@ -70,10 +73,13 @@ custom_kernel_dot_module = Extension(
     include_dirs=paddle_custom_kernel_include,
     library_dirs=paddle_custom_kernel_library_dir,
     libraries=libs,
-    extra_compile_args=paddle_extra_compile_args)
+    extra_compile_args=paddle_extra_compile_args,
+)
 
-setup(name='custom_kernel_dot',
-      version='1.0',
-      description='custom kernel fot compiling',
-      cmdclass={'build_ext': BuildExt},
-      ext_modules=[custom_kernel_dot_module])
+setup(
+    name='custom_kernel_dot',
+    version='1.0',
+    description='custom kernel fot compiling',
+    cmdclass={'build_ext': BuildExt},
+    ext_modules=[custom_kernel_dot_module],
+)