Fix libtensorflow CUDA compute capabilities.

PiperOrigin-RevId: 330987201
Change-Id: I1a10343229216f70da922c163cc9ebcbfe069947
This commit is contained in:
Amit Patankar 2020-09-10 12:05:11 -07:00 committed by TensorFlower Gardener
parent 082ca0493e
commit 660f6e0610
2 changed files with 1 additions and 1 deletions

View File

@ -56,6 +56,7 @@ function build_libtensorflow_tarball() {
if [ "${TF_NEED_CUDA}" == "1" ]; then
BAZEL_OPTS="${BAZEL_OPTS} --config=cuda --crosstool_top=//third_party/toolchains/preconfig/ubuntu16.04/gcc7_manylinux2010-nvcc-cuda11:toolchain"
export TF_NEED_ROCM=0
export TF_CUDA_COMPUTE_CAPABILITIES="sm_35,sm_50,sm_60,sm_70,sm_75,compute_80"
fi
bazel clean --expunge
yes "" | ./configure

View File

@ -19,5 +19,4 @@
set -ex
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
export TF_NEED_CUDA=1
export TF_CUDA_COMPUTE_CAPABILITIES="sm_35,sm_50,sm_60,sm_70,sm_75,compute_80"
"${SCRIPT_DIR}/libtensorflow_docker.sh"