Update AMD Docker image (#35804)

huggingface · Jan 21, 2025 · 8e4cedd · 8e4cedd
1 parent 705aeaa
commit 8e4cedd
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/docker/transformers-pytorch-amd-gpu/Dockerfile b/docker/transformers-pytorch-amd-gpu/Dockerfile
@@ -1,5 +1,4 @@
-FROM rocm/dev-ubuntu-22.04:6.1
-# rocm/pytorch has no version with 2.1.0
+FROM rocm/dev-ubuntu-22.04:6.3
 LABEL maintainer="Hugging Face"
 
 ARG DEBIAN_FRONTEND=noninteractive
@@ -9,9 +8,11 @@ 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 torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm6.1
+RUN python3 -m pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/rocm6.3/
 
 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"