fix: install CU11 PyTorch in the CU11 docker image (deepmodeling#3198)

The default one from PyPI is for CU12. --------- Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
iProzd · Jan 29, 2024 · 4a29c8c · 4a29c8c
1 parent 8eadd3e
commit 4a29c8c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/source/install/docker/Dockerfile b/source/install/docker/Dockerfile
@@ -6,7 +6,8 @@ RUN python -m venv /opt/deepmd-kit
 ENV PATH="/opt/deepmd-kit/bin:$PATH"
 # Install package
 COPY dist /dist
-RUN pip install "$(ls /dist/deepmd_kit${VARIANT}-*manylinux*_x86_64.whl)[gpu,cu${CUDA_VERSION},lmp,ipi,torch]" \
+RUN if [ "${CUDA_VERSION}" = 11 ]; then pip install torch --index-url https://download.pytorch.org/whl/cu118; fi \
+    && pip install "$(ls /dist/deepmd_kit${VARIANT}-*manylinux*_x86_64.whl)[gpu,cu${CUDA_VERSION},lmp,ipi,torch]" \
     && dp -h \
     && lmp -h \
     && dp_ipi \