We may need to install CUDNN.

coreweave · Aug 21, 2024 · 55c9e59 · 55c9e59
1 parent f137049
commit 55c9e59
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/torch-extras/Dockerfile b/torch-extras/Dockerfile
@@ -54,6 +54,7 @@ RUN export \
       libcusolver-dev-${CUDA_PACKAGE_VERSION} \
       cuda-nvprof-${CUDA_PACKAGE_VERSION} \
       cuda-profiler-api-${CUDA_PACKAGE_VERSION} \
+      cudnn9-${CUDA_MAJOR_VERSION} \
       libaio-dev \
       ninja-build && \
     apt-get clean