diff --git a/cmake/cudnn.cmake b/cmake/cudnn.cmake index 9eebea816cbfc91052c95ecf99ecc4b0bea4e4c2..c3c1777e3959362c84d3c8fd4b39156a293304f2 100644 --- a/cmake/cudnn.cmake +++ b/cmake/cudnn.cmake @@ -25,13 +25,30 @@ list(APPEND CUDNN_CHECK_LIBRARY_DIRS $ENV{CUDNN_ROOT} $ENV{CUDNN_ROOT}/lib64 $ENV{CUDNN_ROOT}/lib - /usr/lib) -find_library(CUDNN_LIBRARY NAMES libcudnn.so libcudnn.dylib # libcudnn_static.a + /usr/lib + ${CUDA_TOOLKIT_ROOT_DIR} + ${CUDA_TOOLKIT_ROOT_DIR}/lib/x64 + ) +set(CUDNN_LIB_NAME "") +if (LINUX) +set(CUDNN_LIB_NAME "libcudnn.so") +endif(LINUX) + +if(WIN32) +# only support cudnn7 +set(CUDNN_LIB_NAME "cudnn.lib" "cudnn64_7.dll") +endif(WIN32) + +if(Apple) +set(CUDNN_LIB_NAME "libcudnn.dylib" "libcudnn.so") +endif(Apple) +find_library(CUDNN_LIBRARY NAMES ${CUDNN_LIB_NAME} # libcudnn_static.a PATHS ${CUDNN_CHECK_LIBRARY_DIRS} ${CUDNN_INCLUDE_DIR} ${__libpath_hist} NO_DEFAULT_PATH DOC "Path to cuDNN library.") - +message("Include Dir" ${CUDNN_INCLUDE_DIR}) +message("Library Dir" ${CUDNN_LIBRARY}) if(CUDNN_INCLUDE_DIR AND CUDNN_LIBRARY) set(CUDNN_FOUND ON) else()