Fix libtensorflow CUDA compute capabilities.
PiperOrigin-RevId: 330987201 Change-Id: I1a10343229216f70da922c163cc9ebcbfe069947
This commit is contained in:
parent
082ca0493e
commit
660f6e0610
@ -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
|
||||
|
@ -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"
|
||||
|
Loading…
x
Reference in New Issue
Block a user