31 lines
758 B
Docker
31 lines
758 B
Docker
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" ]
|