diff --git a/pytorch.dockerfile b/pytorch.dockerfile
index 4af61a6..22f9fbd 100644
--- a/pytorch.dockerfile
+++ b/pytorch.dockerfile
@@ -48,9 +48,6 @@ PYTORCH_BUILD_VERSION=$(echo "${PYTORCH_BRANCH}" | tr -d "v") python${PYTHON_VER
 # https://github.com/pytorch/pytorch/issues/125230
 # https://github.com/AngryLoki/pytorch/commit/38d6f115c65b67946de6765165ece0288c8416d4
 
-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
@@ -59,7 +56,7 @@ sudo apt-get clean
 sudo rm -rf /var/lib/apt/lists/*
 EOF
 
-RUN python${PYTHON_VERSION} -m pip install /home/rocm/torch.whl
+RUN python${PYTHON_VERSION} -m pip install /home/rocm/pytorch/dist/*.whl
 
 # Torch Vision
 WORKDIR /home/rocm/
@@ -68,9 +65,6 @@ 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
@@ -79,9 +73,6 @@ 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