From d8f33e3c2bad6e2d024ea58dc6c06af2fae27803 Mon Sep 17 00:00:00 2001
From: Felix Marty <felix@hf.co>
Date: Thu, 25 Jan 2024 12:05:35 +0000
Subject: [PATCH] update torch

---
 Dockerfile_amd | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/Dockerfile_amd b/Dockerfile_amd
index e3aba45e..b6752c9a 100644
--- a/Dockerfile_amd
+++ b/Dockerfile_amd
@@ -75,8 +75,8 @@ RUN chmod +x ~/mambaforge.sh && \
     mamba init && \
     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.
-RUN pip install --pre torch==2.2.0.dev20231106 --index-url https://download.pytorch.org/whl/nightly/rocm5.7
+# Install PyTorch 2.2 RC compiled against RoCm 5.7, as VLLM can not be compiled with RoCm 5.6.
+RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/test/rocm5.7/
 
 FROM base AS kernel-builder