From 9f72a2cce38712222061e61a6e19485a1710d8f3 Mon Sep 17 00:00:00 2001
From: Amit Patankar <amitpatankar@google.com>
Date: Thu, 10 Sep 2020 12:05:11 -0700
Subject: [PATCH] Fix libtensorflow CUDA compute capabilities.

PiperOrigin-RevId: 330987201
Change-Id: I1a10343229216f70da922c163cc9ebcbfe069947
---
 tensorflow/tools/ci_build/builds/libtensorflow.sh    | 1 +
 tensorflow/tools/ci_build/linux/libtensorflow_gpu.sh | 1 -
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/tensorflow/tools/ci_build/builds/libtensorflow.sh b/tensorflow/tools/ci_build/builds/libtensorflow.sh
index a281afe7442..d0be3284eb4 100755
--- a/tensorflow/tools/ci_build/builds/libtensorflow.sh
+++ b/tensorflow/tools/ci_build/builds/libtensorflow.sh
@@ -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-cuda10.1: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
diff --git a/tensorflow/tools/ci_build/linux/libtensorflow_gpu.sh b/tensorflow/tools/ci_build/linux/libtensorflow_gpu.sh
index 99397452acf..6dca0c37c87 100755
--- a/tensorflow/tools/ci_build/linux/libtensorflow_gpu.sh
+++ b/tensorflow/tools/ci_build/linux/libtensorflow_gpu.sh
@@ -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"