[DOCKER] CUDA upgrade to 9.0 to acommodate tensorflow-gpu (1.10.0). (a…

…pache#1761)
zxy844288792 · Sep 25, 2018 · e8d6d9a · e8d6d9a
1 parent 1de91d0
commit e8d6d9a
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/docker/Dockerfile.ci_gpu b/docker/Dockerfile.ci_gpu
@@ -1,5 +1,5 @@
 # CI docker GPU env
-FROM nvidia/cuda:8.0-cudnn7-devel
+FROM nvidia/cuda:9.0-cudnn7-devel
 
 # Base scripts
 RUN apt-get update --fix-missing
@@ -62,6 +62,9 @@ RUN pip3 install Pillow
 COPY install/ubuntu_install_vulkan.sh /install/ubuntu_install_vulkan.sh
 RUN bash /install/ubuntu_install_vulkan.sh
 
+COPY install/ubuntu_install_tensorflow.sh /install/ubuntu_install_tensorflow.sh
+RUN bash /install/ubuntu_install_tensorflow.sh
+
 # AutoTVM deps
 COPY install/ubuntu_install_redis.sh /install/ubuntu_install_redis.sh
 RUN bash /install/ubuntu_install_redis.sh

diff --git a/docker/Dockerfile.demo_gpu b/docker/Dockerfile.demo_gpu
@@ -1,6 +1,6 @@
 # Minimum docker image for demo purposes
 # prebuilt-image: tvmai/demo-gpu
-FROM nvidia/cuda:8.0-cudnn7-devel
+FROM nvidia/cuda:9.0-cudnn7-devel
 
 RUN apt-get update --fix-missing
 

diff --git a/docker/install/ubuntu_install_tensorflow.sh b/docker/install/ubuntu_install_tensorflow.sh
@@ -0,0 +1 @@
+pip3 install tensorflow-gpu