diff --git a/python/setup.py.in b/python/setup.py.in
index 6b0eecb21d45aa5b79a7e55a732ecc28c1e2034a..35033683a80766bd1552681ba4bf10e0526c1f37 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -271,9 +271,11 @@ def find_files(pattern, root):
 
 headers = (
     list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/fluid/framework')) +
+    list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/fluid/imperative')) +
     list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/fluid/memory')) +
     list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/fluid/platform')) +
     list(find_files('*.h', '@PADDLE_SOURCE_DIR@/paddle/fluid/string')) +
+    list(find_files('*.pb.h', '${PADDLE_BINARY_DIR}/paddle/fluid/platform')) +
     list(find_files('*.pb.h', '${PADDLE_BINARY_DIR}/paddle/fluid/framework')) +
     ['${EIGEN_INCLUDE_DIR}/Eigen/Core'] + # eigen
     list(find_files('*', '${EIGEN_INCLUDE_DIR}/Eigen/src')) + # eigen
@@ -286,6 +288,9 @@ headers = (
     list(find_files('*', '${DLPACK_INCLUDE_DIR}')) + # dlpack 
     list(find_files('*.h', '${THREADPOOL_INCLUDE_DIR}'))) # threadpool
 
+if '${WITH_MKLDNN}' == 'ON':
+    headers += list(find_files('*', '${MKLDNN_INSTALL_DIR}/include')) # mkldnn
+
 
 class InstallCommand(InstallCommandBase):
     def finalize_options(self):
@@ -333,6 +338,7 @@ class InstallHeaders(Command):
                        'install/protobuf/include',
                        'install/gflags/include',
                        'install/glog/include', 'install/xxhash/include',
+                       'install/mkldnn/include',
                        'threadpool/src/extern_threadpool']
             for pattern in patterns:
                 install_dir = re.sub(pattern, '', install_dir)