From 47d7d83b6fdf66c3a4aec5b08266992358aa6e8f Mon Sep 17 00:00:00 2001
From: Ian Rodney <ian.rodney@gmail.com>
Date: Thu, 17 Sep 2020 10:56:58 -0700
Subject: [PATCH] [docker] Fix GPU support for tensorflow (#10779)

---
 ci/travis/build-docker-images.sh  | 2 +-
 python/requirements_ml_docker.txt | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/ci/travis/build-docker-images.sh b/ci/travis/build-docker-images.sh
index 6f3eedc68..5a1aa7655 100755
--- a/ci/travis/build-docker-images.sh
+++ b/ci/travis/build-docker-images.sh
@@ -22,7 +22,7 @@ build_and_push_tags() {
     # $2 tag for image (e.g. hash of commit)
     for GPU in "" "-gpu" 
     do 
-        BASE_IMAGE=$(if [ "$GPU" ]; then echo "nvidia/cuda:11.0-cudnn8-runtime-ubuntu18.04"; else echo "ubuntu:focal"; fi;)
+        BASE_IMAGE=$(if [ "$GPU" ]; then echo "nvidia/cuda:10.1-cudnn8-runtime-ubuntu18.04"; else echo "ubuntu:focal"; fi;)
         FULL_NAME_WITH_TAG="rayproject/$1:$2$GPU"
         NIGHTLY_FULL_NAME_WITH_TAG="rayproject/$1:nightly$GPU"
         docker build --no-cache --build-arg GPU="$GPU" --build-arg BASE_IMAGE="$BASE_IMAGE" --build-arg WHEEL_PATH=".whl/$WHEEL" --label "SHA=$2" -t "$FULL_NAME_WITH_TAG" /"$ROOT_DIR"/docker/"$1"
diff --git a/python/requirements_ml_docker.txt b/python/requirements_ml_docker.txt
index ec20ceb0c..2a9cc529a 100644
--- a/python/requirements_ml_docker.txt
+++ b/python/requirements_ml_docker.txt
@@ -1,3 +1,4 @@
 -r requirements_rllib.txt
 -r requirements_tune.txt
-ipython
\ No newline at end of file
+ipython
+tensorflow-gpu
\ No newline at end of file