Add support for V100 GPU. (Kaggle#940)

rosbo · web-flow · commit 91aeeebc2fd5 · 2021-01-22T11:27:03.000-08:00
TensorFlow must be built with 7.0 compute capability.

http://b/178104927
diff --git a/tensorflow-whl/CHANGELOG.md b/tensorflow-whl/CHANGELOG.md
@@ -20,4 +20,5 @@
 * `2.3.1-py37`: TensorFlow 2.3.1 with Python 3.7
 * `2.3.1-py37-2`: TensorFlow 2.3.1 & TFA 0.11 with Python 3.7
 * `2.4.0-py37`: TensorFlow 2.4.0 & Python 3.7 & TFA 0.12
-* `2.4.0-py37-2`: TensorFlow 2.4.0 & m61 base image
+* `2.4.0-py37-2`: TensorFlow 2.4.0 & m61 base image
+* `2.4.0-py37-3`: TensorFlow 2.4.0 & 7.0 CUDA compute capability
diff --git a/tensorflow-whl/Dockerfile b/tensorflow-whl/Dockerfile
@@ -103,8 +103,8 @@ RUN cd /usr/local/src/tensorflow_gcs_config && \
 # Create a tensorflow wheel for GPU/cuda
 ENV TF_NEED_CUDA=1
 ENV TF_CUDA_VERSION=$CUDA_MAJOR_VERSION.$CUDA_MINOR_VERSION
-# 3.7 is for the K80 and 6.0 is for the P100, 7.5 is for the T4: https://developer.nvidia.com/cuda-gpus
-ENV TF_CUDA_COMPUTE_CAPABILITIES=3.7,6.0,7.5
+# 3.7 (K80), 6.0 (P100), 7.0 (V100), 7.5 (T4): https://developer.nvidia.com/cuda-gpus
+ENV TF_CUDA_COMPUTE_CAPABILITIES=3.7,6.0,7.0,7.5
 ENV TF_CUDNN_VERSION=7
 ENV TF_NCCL_VERSION=2
 ENV NCCL_INSTALL_PATH=/usr/
