Default CUDNN_HOME to CUDA_TOOLKIT_TARGET_DIR. The cuDNN distro is most naturally installed in the same directory as the CUDA SDK, so try to find it there if the user doesn't specify any other directory.

PiperOrigin-RevId: 157436253
This commit is contained in:
A. Unique TensorFlower 2017-05-30 00:48:35 -07:00 committed by TensorFlower Gardener
parent eb7cf93315
commit 73d10599fe

View File

@ -200,6 +200,9 @@ if (tensorflow_ENABLE_GPU)
add_definitions(-DGOOGLE_CUDA=1 -DTF_EXTRA_CUDA_CAPABILITIES=3.0,3.5,5.2)
# add cudnn
if(NOT CUDNN_HOME)
set(CUDNN_HOME ${CUDA_TOOLKIT_TARGET_DIR})
endif(NOT CUDNN_HOME)
include_directories(${CUDNN_HOME})
set(CUDA_LIBRARIES ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUBLAS_LIBRARIES} ${CUDA_CUFFT_LIBRARIES}
${CUDA_curand_LIBRARY} ${CUDA_cupti_LIBRARY} ${CUDA_cusolver_LIBRARY} ${CUDNN_HOME}/lib/x64/cudnn.lib)