update torch

This commit is contained in:
Felix Marty 2024-01-25 12:05:35 +00:00 committed by OlivierDehaene
parent 33111a0e7f
commit d8f33e3c2b

View File

@ -75,8 +75,8 @@ RUN chmod +x ~/mambaforge.sh && \
mamba init && \ mamba init && \
rm ~/mambaforge.sh rm ~/mambaforge.sh
# Install PyTorch nightly (2.2.0.dev2023) compiled against RoCm 5.7, as VLLM can not be compiled with RoCm 5.6. # Install PyTorch 2.2 RC compiled against RoCm 5.7, as VLLM can not be compiled with RoCm 5.6.
RUN pip install --pre torch==2.2.0.dev20231106 --index-url https://download.pytorch.org/whl/nightly/rocm5.7 RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/test/rocm5.7/
FROM base AS kernel-builder FROM base AS kernel-builder