update rocm version

huggingface · Aug 30, 2024 · 0ad52b3 · 0ad52b3
1 parent ca841a7
commit 0ad52b3
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/docker/rocm/Dockerfile b/docker/rocm/Dockerfile
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-ARG ROCM_VERSION=5.7.1
+ARG ROCM_VERSION=6.1.2
 ARG UBUNTU_VERSION=22.04
 
 FROM rocm/dev-ubuntu-${UBUNTU_VERSION}:${ROCM_VERSION}
@@ -29,8 +29,8 @@ RUN apt-get update && apt-get upgrade -y && apt-get install -y --no-install-reco
     cd /opt/rocm/share/amd_smi && pip install .
 
 # Install PyTorch
-ARG TORCH_ROCM=rocm5.7
-ARG TORCH_VERSION=2.2.2
+ARG TORCH_ROCM=rocm6.1
+ARG TORCH_VERSION=stable
 
 RUN if [ "${TORCH_VERSION}" = "stable" ]; then \
     pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/${TORCH_ROCM} ; \