From d1a82edbaf1555be6778a904b4929b8a865db102 Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Tue, 10 Sep 2024 15:26:27 -0700
Subject: [PATCH] feat: python version

---
 pytorch.dockerfile | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 78e60af..6efa626 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -17,14 +17,21 @@ RUN git submodule update --init --recursive
 
 RUN <<EOF
 sudo apt-get update
-sudo apt install -y python3-dev
+sudo apt install -y software-properties-common
+sudo add-apt-repository -y ppa:deadsnakes/ppa
+sudo apt-get update
+sudo apt-get install -y python${PYTHON_VERSION} python${PYTHON_VERSION}-dev
 sudo apt-get clean
 sudo rm -rf /var/lib/apt/lists/*
 EOF
 
-RUN pip install typing-extensions
+RUN python${PYTHON_VERSION} -m pip install typing-extensions wheel
 
 ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
-RUN python3 tools/amd_build/build_amd.py && USE_ROCM=1 python3 setup.py bdist_wheel
+RUN python${PYTHON_VERSION} tools/amd_build/build_amd.py && USE_ROCM=1 python${PYTHON_VERSION} setup.py bdist_wheel
 
+WORKDIR /home/rocm/pytorch/dist
+RUN cp (ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/pytorch-rocm.whl
+
+WORKDIR /home/rocm
 CMD [ "bash", "-l" ]