wtongze/rocm
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: torch vision and audio version

Browse Source
This commit is contained in:
wtongze
2024-09-12 13:29:28 -07:00
parent f989495458
commit 7946c9eaab
1 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch.dockerfile
+8 -3
View File
@@ -59,10 +59,13 @@ RUN python${PYTHON_VERSION} -m pip install /home/rocm/pytorch/dist/*.whl
# Torch Vision
WORKDIR /home/rocm/
RUN git clone --depth 1 https://github.com/pytorch/vision.git -b v0.17.2 vision
ENV TORCHVISION_BRANCH="v0.17.2"
RUN git clone --depth 1 https://github.com/pytorch/vision.git -b ${TORCHVISION_BRANCH} vision
WORKDIR /home/rocm/vision
RUN USE_ROCM=ON python${PYTHON_VERSION} setup.py bdist_wheel
RUN USE_ROCM=ON \
BUILD_VERSION=$(echo "${TORCHVISION_BRANCH}" | tr -d "v") \
python${PYTHON_VERSION} setup.py bdist_wheel
# Torch Audio
WORKDIR /home/rocm/
@@ -70,7 +73,7 @@ RUN git clone --depth 1 https://github.com/pytorch/audio.git -b ${PYTORCH_BRANCH
WORKDIR /home/rocm/audio
RUN USE_ROCM=ON \
PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
python${PYTHON_VERSION} setup.py bdist_wheel
RUN <<EOF
@@ -80,5 +83,7 @@ cp /home/rocm/vision/dist/*.whl /home/rocm/build
cp /home/rocm/audio/dist/*.whl /home/rocm/build
EOF
RUN python${PYTHON_VERSION} -m pip install /home/rocm/build/*.whl
WORKDIR /home/rocm
CMD [ "bash", "-l" ]