feat: basic pytorch

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