diff --git a/python/setup.py.in b/python/setup.py.in
index 748880302f274240ca73662bdb7d249471f0a5fd..c28aaa1ee38a6216c701f3d279bdf6baabec4f55 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -286,11 +286,13 @@ 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
 
+if '${WITH_MKLDNN}' == 'ON':
+    headers += list(find_files('*', '${MKLDNN_INSTALL_DIR}/include')) # mkldnn
+
 
 class InstallCommand(InstallCommandBase):
     def finalize_options(self):