From 4fc0937174cc47d052e6dd27fa6da043adcb2bc2 Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Thu, 12 Sep 2024 13:48:03 -0700
Subject: [PATCH] fix: numpy in torch

---
 pytorch.dockerfile | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 77870b6..ccdfd93 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -4,6 +4,7 @@ FROM wtongze/rocm-base:${ROCM_VERSION}
 
 ARG ROCM_VERSION="5.7.3"
 ARG PYTORCH_BRANCH="v2.2.2"
+ARG TORCHVISION_BRANCH="v0.17.2"
 ARG PYTHON_VERSION="3.11"
 ARG HIP_ARCH="gfx803;gfx900"
 ARG DEBIAN_FRONTEND="noninteractive"
@@ -26,7 +27,8 @@ sudo apt-get clean
 sudo rm -rf /var/lib/apt/lists/*
 EOF
 
-RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel
+RUN python${PYTHON_VERSION} -m pip install wheel
+RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
 
 RUN <<EOF
 sudo apt-get update
@@ -38,6 +40,10 @@ EOF
 ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
 ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
 ENV PYTORCH_BUILD_NUMBER=0
+
+# Torch
+# > https://github.com/pytorch/pytorch/issues/125230
+# >https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
 RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \
 USE_ROCM=ON \
 USE_NUMPY=1 \
@@ -45,9 +51,6 @@ USE_FLASH_ATTENTION=OFF \
 USE_MEM_EFF_ATTENTION=OFF \
 PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") python${PYTHON_VERSION} setup.py bdist_wheel
 
-# https://github.com/pytorch/pytorch/issues/125230
-# https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
-
 RUN <<EOF
 sudo apt-get update
 sudo apt-get install -y ffmpeg libjpeg-dev libpng-dev
@@ -59,7 +62,6 @@ RUN python${PYTHON_VERSION} -m pip install /home/rocm/pytorch/dist/*.whl
 
 # Torch Vision
 WORKDIR /home/rocm/
-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
@@ -72,6 +74,7 @@ WORKDIR /home/rocm/
 RUN git clone --depth 1 https://github.com/pytorch/audio.git -b ${PYTORCH_BRANCH} audio
 
 WORKDIR /home/rocm/audio
+RUN python${PYTHON_VERSION} -m pip install -r requirements.txt
 RUN USE_ROCM=ON \
 BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
 python${PYTHON_VERSION} setup.py bdist_wheel