fix: build config and depenency

pytorch.dockerfile

@@ -2,7 +2,9 @@ ARG ROCM_VERSION="5.7.3"
 
 FROM wtongze/rocm-base:${ROCM_VERSION}
 
-ARG PYTORCH_BRANCH="v2.4.1"
+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"
@@ -27,11 +29,33 @@ EOF
 
 RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel
 
+# For ROCm 4.5+
+RUN <<EOF
+sudo apt-get update
+sudo apt-get install -y rocm-ml-sdk rocm-hip-sdk
+sudo apt-get clean
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
+RUN <<EOF
+sudo apt-get update
+sudo apt-get install -y ninja-build
+sudo apt-get clean
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
+ENV ROCM_PATH=/opt/rocm-${ROCM_VERSION}
 ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
-RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && USE_ROCM=1 python${PYTHON_VERSION} setup.py bdist_wheel
+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_FLASH_ATTENTION=OFF USE_MEM_EFF_ATTENTION=OFF python${PYTHON_VERSION} setup.py bdist_wheel
+
+# https://github.com/pytorch/pytorch/issues/125230
+# https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
 
 WORKDIR /home/rocm/pytorch/dist
-RUN cp (ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/pytorch-rocm.whl
+RUN cp $(ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/pytorch-rocm.whl
 
 WORKDIR /home/rocm
 CMD [ "bash", "-l" ]