hiro/ray
1
0
Fork 0
mirror of https://github.com/vale981/ray synced 2025-03-06 10:31:39 -05:00
Code Issues Projects Releases Packages Wiki Activity Actions

CUDA 11.2 in docker images

Browse source
This commit is contained in:
Sven Mika 2021-08-16 12:31:19 +02:00 • committed by Kai Fricke
parent 34de954723
commit ee9520d4e1
2 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
ci/travis/build-docker-images.py
View file

@ -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"

2
python/requirements_ml_docker.txt
View file

@ -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