fix pytorch and cuda

js-ish · May 20, 2024 · 4b5d31f · 4b5d31f
1 parent 53ae90c
commit 4b5d31f
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/pytorch/base.Dockerfile b/pytorch/base.Dockerfile
@@ -1,4 +1,4 @@
-FROM nvidia/cuda:12.1.0-devel-ubuntu20.04
+FROM nvidia/cuda:12.4.1-devel-ubuntu22.04
 
 RUN set -eux; \
     apt-get update; \
@@ -11,6 +11,6 @@ RUN set -eux; \
     rm -rf /var/lib/apt/lists/*
 
 RUN pip install --disable-pip-version-check --no-cache-dir \
-       torch==2.1.0 torchvision==0.16 torchaudio==2.1.0 torchtext==0.16.0; \
+        torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu121; \
     pip install --disable-pip-version-check --no-cache-dir \
-       Pillow scikit-learn pyyaml
+        Pillow scikit-learn pyyaml