From ee9520d4e1a65823cf2407cb2ea31710851c8511 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Mon, 16 Aug 2021 12:31:19 +0200 Subject: [PATCH] CUDA 11.2 in docker images --- ci/travis/build-docker-images.py | 2 +- python/requirements_ml_docker.txt | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/ci/travis/build-docker-images.py b/ci/travis/build-docker-images.py index e5ae97b90..774ca19f0 100644 --- a/ci/travis/build-docker-images.py +++ b/ci/travis/build-docker-images.py @@ -133,7 +133,7 @@ def _build_cpu_gpu_images(image_name, no_cache=True) -> List[str]: if image_name == "base-deps": build_args["BASE_IMAGE"] = ( - "nvidia/cuda:11.0-cudnn8-devel-ubuntu18.04" + "nvidia/cuda:11.2.0-cudnn8-devel-ubuntu18.04" if gpu == "-gpu" else "ubuntu:focal") else: # NOTE(ilr) This is a bit of an abuse of the name "GPU" diff --git a/python/requirements_ml_docker.txt b/python/requirements_ml_docker.txt index 8d07653ae..e22fb024a 100644 --- a/python/requirements_ml_docker.txt +++ b/python/requirements_ml_docker.txt @@ -1,6 +1,6 @@ ipython # In TF >v2, GPU support is included in the base package. -tensorflow==2.4.1 +tensorflow==2.5 -f https://download.pytorch.org/whl/torch_stable.html torch==1.8.1+cu111 -f https://download.pytorch.org/whl/torch_stable.html