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" ]