diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 295ab5d..77870b6 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -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" ]