From b55c35919473a459dc83aa08083be835e2c93754 Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Thu, 12 Sep 2024 12:32:37 -0700
Subject: [PATCH] feat: simplify pytorch version

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

diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 7d73297..c552d2f 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -4,7 +4,6 @@ FROM wtongze/rocm-base:${ROCM_VERSION}
 
 ARG ROCM_VERSION="5.7.3"
 ARG PYTORCH_BRANCH="v2.2.2"
-ARG PYTORCH_BUILD_VERSION="2.2.2"
 ARG PYTHON_VERSION="3.11"
 ARG HIP_ARCH="gfx803;gfx900"
 ARG DEBIAN_FRONTEND="noninteractive"
@@ -38,10 +37,13 @@ EOF
 
 ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
 ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
-ENV PYTORCH_BUILD_VERSION=${PYTORCH_BUILD_VERSION}
 ENV PYTORCH_BUILD_NUMBER=0
 RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && \
-USE_ROCM=ON USE_NUMPY=1 USE_FLASH_ATTENTION=OFF USE_MEM_EFF_ATTENTION=OFF python${PYTHON_VERSION} setup.py bdist_wheel
+USE_ROCM=ON \
+USE_NUMPY=1 \
+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