Add cudnn to build.

coreweave · Aug 28, 2024 · 422dba3 · 422dba3
1 parent 6283825
commit 422dba3
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torch/Dockerfile b/torch/Dockerfile
@@ -120,6 +120,7 @@ RUN export \
       CUDA_PACKAGE_VERSION="${CUDA_MAJOR_VERSION}-${CUDA_MINOR_VERSION}" && \
     apt-get -qq update && \
     apt-get -qq install --no-upgrade -y \
+        cudnn9-cuda-${CUDA_MAJOR_VERSION} \
 	cuda-nvtx-${CUDA_PACKAGE_VERSION} \
 	cuda-nvrtc-dev-${CUDA_PACKAGE_VERSION}