Update Dockerfile.cuda

doppeltilde · Jun 13, 2024 · 5cae01a · 5cae01a
1 parent e9ce750
commit 5cae01a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/Dockerfile.cuda b/Dockerfile.cuda
@@ -1,8 +1,8 @@
-FROM nvidia/cuda:12.2.2-cudnn8-runtime-ubuntu22.04
+FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04
 WORKDIR /app
 COPY . /app
 RUN apt-get update && \
     apt-get install -y python3-pip && \
     pip install --upgrade pip
-RUN pip install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/121
+RUN pip install --no-cache-dir -r requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121
 CMD ["fastapi", "run", "main.py", "--proxy-headers", "--host", "0.0.0.0", "--port", "8000"]