From 7946c9eaab8eb5d2716362ec5e225e79d8bb1b9b Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Thu, 12 Sep 2024 13:29:28 -0700
Subject: [PATCH] fix: torch vision and audio version

---
 pytorch.dockerfile | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 295ab5d..77870b6 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -59,10 +59,13 @@ RUN python${PYTHON_VERSION} -m pip install /home/rocm/pytorch/dist/*.whl
 
 # Torch Vision
 WORKDIR /home/rocm/
-RUN git clone --depth 1 https://github.com/pytorch/vision.git -b v0.17.2 vision
+ENV TORCHVISION_BRANCH="v0.17.2"
+RUN git clone --depth 1 https://github.com/pytorch/vision.git -b ${TORCHVISION_BRANCH} vision
 
 WORKDIR /home/rocm/vision
-RUN USE_ROCM=ON python${PYTHON_VERSION} setup.py bdist_wheel
+RUN USE_ROCM=ON \
+BUILD_VERSION=$(echo "${TORCHVISION_BRANCH}" | tr -d "v") \
+python${PYTHON_VERSION} setup.py bdist_wheel
 
 # Torch Audio
 WORKDIR /home/rocm/
@@ -70,7 +73,7 @@ RUN git clone --depth 1 https://github.com/pytorch/audio.git -b ${PYTORCH_BRANCH
 
 WORKDIR /home/rocm/audio
 RUN USE_ROCM=ON \
-PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
+BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
 python${PYTHON_VERSION} setup.py bdist_wheel
 
 RUN <<EOF
@@ -80,5 +83,7 @@ cp /home/rocm/vision/dist/*.whl /home/rocm/build
 cp /home/rocm/audio/dist/*.whl /home/rocm/build
 EOF
 
+RUN python${PYTHON_VERSION} -m pip install /home/rocm/build/*.whl
+
 WORKDIR /home/rocm
 CMD [ "bash", "-l" ]