From 96625d85fd3476d716dbf3d6cde71c982bb0a239 Mon Sep 17 00:00:00 2001
From: ivarflakstad <69173633+ivarflakstad@users.noreply.github.com>
Date: Tue, 28 Jan 2025 11:10:28 +0100
Subject: [PATCH] Use rocm6.2 for AMD images (#35930)

* Use rocm6.2 as rocm6.3 only has nightly pytorch wheels atm

* Use stable wheel index for torch libs
---
 docker/transformers-pytorch-amd-gpu/Dockerfile           | 6 ++----
 docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile | 6 +++---
 2 files changed, 5 insertions(+), 7 deletions(-)

diff --git a/docker/transformers-pytorch-amd-gpu/Dockerfile b/docker/transformers-pytorch-amd-gpu/Dockerfile
index 88763ca07f4..df03cb041a1 100644
--- a/docker/transformers-pytorch-amd-gpu/Dockerfile
+++ b/docker/transformers-pytorch-amd-gpu/Dockerfile
@@ -1,4 +1,4 @@
-FROM rocm/dev-ubuntu-22.04:6.3
+FROM rocm/dev-ubuntu-22.04:6.2.4
 LABEL maintainer="Hugging Face"
 
 ARG DEBIAN_FRONTEND=noninteractive
@@ -8,11 +8,9 @@ RUN apt update && \
     apt clean && \
     rm -rf /var/lib/apt/lists/*
 
-RUN export PATH="${PATH:+${PATH}:}~/opt/rocm/bin"
-
 RUN python3 -m pip install --no-cache-dir --upgrade pip numpy
 
-RUN python3 -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
+RUN python3 -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.2
 
 RUN python3 -m pip install --no-cache-dir --upgrade importlib-metadata setuptools ninja git+https://github.com/facebookresearch/detectron2.git pytesseract "itsdangerous<2.1.0"
 
diff --git a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
index b67793dc959..a8edb8ff03e 100644
--- a/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
+++ b/docker/transformers-pytorch-deepspeed-amd-gpu/Dockerfile
@@ -1,11 +1,11 @@
-FROM rocm/dev-ubuntu-22.04:6.3
+FROM rocm/dev-ubuntu-22.04:6.2.4
 LABEL maintainer="Hugging Face"
 
 ARG DEBIAN_FRONTEND=noninteractive
 ARG PYTORCH='2.5.1'
 ARG TORCH_VISION='0.20.0'
 ARG TORCH_AUDIO='2.5.0'
-ARG ROCM='6.3'
+ARG ROCM='6.2'
 
 RUN apt update && \
     apt install -y --no-install-recommends \
@@ -45,4 +45,4 @@ RUN cd transformers && python3 setup.py develop
 RUN python3 -c "from deepspeed.launcher.runner import main"
 
 # Remove nvml as it is not compatible with ROCm
-RUN python3 -m pip uninstall py3nvml pynvml -y
+RUN python3 -m pip uninstall py3nvml pynvml nvidia-ml-py apex -y