diff --git a/python/setup.py.in b/python/setup.py.in
index 7a9730e3e81739e4faaedd132e90109bd504c7c1..748880302f274240ca73662bdb7d249471f0a5fd 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -273,9 +273,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
@@ -284,6 +286,7 @@ headers = (
     list(find_files('*', '${GLOG_INSTALL_DIR}/include')) + # glog
     list(find_files('*', '${BOOST_INCLUDE_DIR}/boost')) + # boost
     list(find_files('*', '${XXHASH_INSTALL_DIR}/include')) + # xxhash
+    list(find_files('mkldnn.h', '${MKLDNN_INSTALL_DIR}/include')) + # mkldnn
     list(find_files('*', '${PROTOBUF_INCLUDE_DIR}')) + # protobuf
     list(find_files('*', '${DLPACK_INCLUDE_DIR}')) + # dlpack 
     list(find_files('*.h', '${THREADPOOL_INCLUDE_DIR}'))) # threadpool
@@ -335,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)