From 1476c395887f8f692f8eef6991eab1542e1bab79 Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Tue, 10 Sep 2024 15:16:52 -0700
Subject: [PATCH] feat: basic pytorch

---
 pytorch.dockerfile | 30 ++++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)
 create mode 100644 pytorch.dockerfile

diff --git a/pytorch.dockerfile b/pytorch.dockerfile
new file mode 100644
index 0000000..78e60af
--- /dev/null
+++ b/pytorch.dockerfile
@@ -0,0 +1,30 @@
+ARG ROCM_VERSION="5.7.3"
+
+FROM wtongze/rocm-base:${ROCM_VERSION}
+
+ARG PYTORCH_BRANCH="v2.4.1"
+ARG PYTHON_VERSION="3.11"
+ARG HIP_ARCH="gfx803;gfx900"
+ARG DEBIAN_FRONTEND="noninteractive"
+
+# Add GPU target
+RUN echo "${HIP_ARCH}" | sed 's/;/\n/g' | sudo tee /opt/rocm/bin/target.lst
+
+WORKDIR /home/rocm
+RUN git clone --depth 1 https://github.com/pytorch/pytorch.git -b ${PYTORCH_BRANCH} pytorch
+WORKDIR /home/rocm/pytorch
+RUN git submodule update --init --recursive
+
+RUN <<EOF
+sudo apt-get update
+sudo apt install -y python3-dev
+sudo apt-get clean
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
+RUN pip install typing-extensions
+
+ENV PYTORCH_ROCM_ARCH=${HIP_ARCH}
+RUN python3 tools/amd_build/build_amd.py && USE_ROCM=1 python3 setup.py bdist_wheel
+
+CMD [ "bash", "-l" ]