diff options
Diffstat (limited to 'tensorflow/contrib/cmake/CMakeLists.txt')
| -rw-r--r-- | tensorflow/contrib/cmake/CMakeLists.txt | 4 |
1 files changed, 3 insertions, 1 deletions
diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt index 64262fdce5..68929da5c9 100644 --- a/tensorflow/contrib/cmake/CMakeLists.txt +++ b/tensorflow/contrib/cmake/CMakeLists.txt @@ -170,7 +170,8 @@ if (tensorflow_ENABLE_GPU) # add cudnn include_directories(${CUDNN_HOME}) - set(CUDA_LIBRARIES ${CUDA_LIBRARIES} ${CUDNN_HOME}/lib/x64/cudnn.lib) + set(CUDA_LIBRARIES ${CUDA_LIBRARIES} ${CUDA_CUDA_LIBRARY} ${CUDA_CUBLAS_LIBRARIES} ${CUDA_CUFFT_LIBRARIES} + ${CUDA_curand_LIBRARY} ${CUDA_cupti_LIBRARY} ${CUDNN_HOME}/lib/x64/cudnn.lib) # create cuda_config.h FILE(WRITE ${tensorflow_source_dir}/third_party/gpus/cuda/cuda_config.h @@ -179,6 +180,7 @@ if (tensorflow_ENABLE_GPU) "#define TF_CUDA_CAPABILITIES CudaVersion(\"3.0\"),CudaVersion(\"3.5\"),CudaVersion(\"5.2\")\n" "#define TF_CUDA_VERSION \"64_80\"\n" "#define TF_CUDNN_VERSION \"64_5\"\n" + "#define TF_CUDA_TOOLKIT_PATH \"${CUDA_TOOLKIT_ROOT_DIR}\"\n" "#endif // CUDA_CUDA_CONFIG_H_\n" ) |