Merge pull request #285 from ahmetoner/fix-toch-cuda

Fix issues with Torch CUDA and cuDNN, fixes #283 and #284
ahmetoner · Feb 17, 2025 · 53f60f3 · 53f60f3
2 parents a8d0dcf + 9bf7d8b
commit 53f60f3
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 451 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ Changelog
 Unreleased
 ----------
 
+### Fixed
+
+- Fixed issues with Torch CUDA and cuDNN
+
 [1.8.0] (2025-02-17)
 --------------------
 

diff --git a/Dockerfile.gpu b/Dockerfile.gpu
@@ -14,6 +14,8 @@ RUN export DEBIAN_FRONTEND=noninteractive \
     python${PYTHON_VERSION} \
     python${PYTHON_VERSION}-venv \
     python3-pip \
+    libcudnn8 \
+    python3-pip \
     && rm -rf /var/lib/apt/lists/*
 
 RUN ln -s -f /usr/bin/python${PYTHON_VERSION} /usr/bin/python3 && \
@@ -39,7 +41,7 @@ COPY --from=swagger-ui /usr/share/nginx/html/swagger-ui.css swagger-ui-assets/sw
 COPY --from=swagger-ui /usr/share/nginx/html/swagger-ui-bundle.js swagger-ui-assets/swagger-ui-bundle.js
 
 RUN poetry install
-RUN $POETRY_VENV/bin/pip install torch==2.6.0 --index-url https://download.pytorch.org/whl/cu121
+RUN $POETRY_VENV/bin/pip install torch==2.6.0+cu126 torchaudio==2.6.0+cu126 --index-url https://download.pytorch.org/whl/cu126
 
 EXPOSE 9000