diff --git a/infra/ansible/e2e_tests.Dockerfile b/infra/ansible/e2e_tests.Dockerfile
index d6c421e7a76..7bedc096c84 100644
--- a/infra/ansible/e2e_tests.Dockerfile
+++ b/infra/ansible/e2e_tests.Dockerfile
@@ -23,6 +23,9 @@ RUN ansible-playbook -vvv playbook.yaml -e "stage=release" -e "${ansible_vars}"
 # Copy test sources.
 RUN mkdir -p /src/pytorch/xla
 COPY --from=build /src/pytorch/xla/test /src/pytorch/xla/test
+# Copy ci_commit_pins from upstream
+RUN mkdir -p /src/pytorch/.github
+COPY --from=build /src/pytorch/.github/ci_commit_pins /src/pytorch/.github/ci_commit_pins
 
 # Copy and install wheels.
 WORKDIR /tmp/wheels
diff --git a/test/tpu/xla_test_job.yaml b/test/tpu/xla_test_job.yaml
index 9691fde766f..4b179fe66a5 100644
--- a/test/tpu/xla_test_job.yaml
+++ b/test/tpu/xla_test_job.yaml
@@ -41,8 +41,8 @@ spec:
     - bash
     - -cxe
     - |
-      # TODO: new TPUCI please remove this short term hack for `RuntimeError: operator torchvision::nms does not exist`
-      pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
+      # Torchvision installed from a commit requires torch, need to install after torch is compiled.
+      pip install --no-use-pep517 "git+https://github.com/pytorch/vision.git@$(cat /src/pytorch/.github/ci_commit_pins/vision.txt)"
       pip install expecttest==0.1.6
       pip install rich
       pip install torch_xla[pallas] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html