feat: vision and audio

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