From cf6e34d46331f9198d4eed8bf57dcc3c46caf87f Mon Sep 17 00:00:00 2001 From: Vincent Roseberry Date: Thu, 21 Jan 2021 18:40:25 +0000 Subject: [PATCH] Add support for V100 GPU. TensorFlow must be built with 7.0 compute capability. http://b/178104927 --- tensorflow-whl/CHANGELOG.md | 3 ++- tensorflow-whl/Dockerfile | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/tensorflow-whl/CHANGELOG.md b/tensorflow-whl/CHANGELOG.md index 804daeee..e0582885 100644 --- 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 \ No newline at end of file +* `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 \ No newline at end of file diff --git a/tensorflow-whl/Dockerfile b/tensorflow-whl/Dockerfile index 7553ea78..48596189 100644 --- 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/