From c9cbb89696c2f5e4e1e8ae644a1987a3e7451577 Mon Sep 17 00:00:00 2001 From: Sean Morgan Date: Thu, 9 Apr 2020 22:31:55 -0400 Subject: [PATCH] * Specify the CUDA information within docker --- tools/docker/gpu_tests.Dockerfile | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tools/docker/gpu_tests.Dockerfile b/tools/docker/gpu_tests.Dockerfile index 6f1f52045f..6c36a8666c 100644 --- a/tools/docker/gpu_tests.Dockerfile +++ b/tools/docker/gpu_tests.Dockerfile @@ -1,5 +1,9 @@ FROM tensorflow/tensorflow:2.1.0-custom-op-gpu-ubuntu16 ENV TF_NEED_CUDA="1" +ENV TF_CUDA_VERSION="10.1" +ENV CUDA_TOOLKIT_PATH="/usr/local/cuda" +ENV TF_CUDNN_VERSION="7" +ENV CUDNN_INSTALL_PATH="/usr/lib/x86_64-linux-gnu" RUN python3 -m pip install --upgrade pip setuptools auditwheel==2.0.0