From c52ca16334dd1ab2eed00ddae1c11bdd5b3a1041 Mon Sep 17 00:00:00 2001
From: Tongze Wang <wtongze@gmail.com>
Date: Thu, 12 Sep 2024 13:02:34 -0700
Subject: [PATCH] feat: vision and audio

---
 pytorch.dockerfile | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 97764c0..534e089 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -51,5 +51,36 @@ PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") python${PYTHON_VER
 WORKDIR /home/rocm/pytorch/dist
 RUN cp $(ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/torch.whl
 
+# Test
+RUN <<EOF
+sudo apt-get update
+sudo apt-get install -y ffmpeg
+sudo apt-get clean
+sudo rm -rf /var/lib/apt/lists/*
+EOF
+
+# Torch Vision
+WORKDIR /home/rocm/
+RUN git clone --depth 1 https://github.com/pytorch/vision.git -b v0.17.2 vision
+
+WORKDIR /home/rocm/vision
+RUN USE_ROCM=ON python${PYTHON_VERSION} setup.py bdist_wheel
+
+WORKDIR /home/rocm/vision/dist
+RUN cp $(ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/torchvision.whl
+
+# Torch Audio
+WORKDIR /home/rocm/
+RUN git clone --depth 1 https://github.com/pytorch/audio.git -b ${PYTORCH_BRANCH} audio
+
+WORKDIR /home/rocm/audio
+RUN USE_ROCM=ON \
+PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") \
+python${PYTHON_VERSION} setup.py bdist_wheel
+
+WORKDIR /home/rocm/audio/dist
+RUN cp $(ls -l *.whl | awk '{ print $9 }' | tail -n1 | tr -d '\n') /home/rocm/torchaudio.whl
+# Test End
+
 WORKDIR /home/rocm
 CMD [ "bash", "-l" ]