aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/cmake
diff options
context:
space:
mode:
Diffstat (limited to 'tensorflow/contrib/cmake')
-rw-r--r--tensorflow/contrib/cmake/CMakeLists.txt4
-rw-r--r--tensorflow/contrib/cmake/tf_cc_ops.cmake3
-rw-r--r--tensorflow/contrib/cmake/tf_python.cmake3
3 files changed, 8 insertions, 2 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"
)
diff --git a/tensorflow/contrib/cmake/tf_cc_ops.cmake b/tensorflow/contrib/cmake/tf_cc_ops.cmake
index 6eaa2502be..bca700aca2 100644
--- a/tensorflow/contrib/cmake/tf_cc_ops.cmake
+++ b/tensorflow/contrib/cmake/tf_cc_ops.cmake
@@ -71,7 +71,7 @@ foreach(tf_cc_op_lib_name ${tf_cc_op_lib_names})
COMMAND ${tf_cc_op_lib_name}_gen_cc ${cc_ops_target_dir}/${tf_cc_op_lib_name}.h ${cc_ops_target_dir}/${tf_cc_op_lib_name}.cc ${tensorflow_source_dir}/tensorflow/cc/ops/op_gen_overrides.pbtxt ${cc_ops_include_internal}
DEPENDS ${tf_cc_op_lib_name}_gen_cc create_cc_ops_header_dir
)
-
+
list(APPEND tf_cc_ops_generated_files ${cc_ops_target_dir}/${tf_cc_op_lib_name}.h)
list(APPEND tf_cc_ops_generated_files ${cc_ops_target_dir}/${tf_cc_op_lib_name}.cc)
list(APPEND tf_cc_ops_generated_files ${cc_ops_target_dir}/${tf_cc_op_lib_name}_internal.h)
@@ -79,6 +79,7 @@ foreach(tf_cc_op_lib_name ${tf_cc_op_lib_names})
endforeach()
+
########################################################
# tf_cc_ops library
########################################################
diff --git a/tensorflow/contrib/cmake/tf_python.cmake b/tensorflow/contrib/cmake/tf_python.cmake
index 7717cf7b71..9ab6f176c7 100644
--- a/tensorflow/contrib/cmake/tf_python.cmake
+++ b/tensorflow/contrib/cmake/tf_python.cmake
@@ -372,6 +372,9 @@ add_python_module("tensorflow/contrib/slim/python/slim/nets")
add_python_module("tensorflow/contrib/solvers")
add_python_module("tensorflow/contrib/solvers/python")
add_python_module("tensorflow/contrib/solvers/python/ops")
+add_python_module("tensorflow/contrib/sparsemax")
+add_python_module("tensorflow/contrib/sparsemax/python")
+add_python_module("tensorflow/contrib/sparsemax/python/ops")
add_python_module("tensorflow/contrib/specs")
add_python_module("tensorflow/contrib/specs/python")
add_python_module("tensorflow/contrib/stat_summarizer")