diff --git a/.github/workflows/linux-test.yml b/.github/workflows/linux-test.yml
index 9f366c30e3..0df8f76d3c 100644
--- a/.github/workflows/linux-test.yml
+++ b/.github/workflows/linux-test.yml
@@ -25,7 +25,7 @@ on:
         type: string
       test-infra-ref:
         description: "Test infra reference to use"
-        default: ""
+        default: "release/2.2"
         type: string
       build-matrix:
         description: "Build matrix to utilize"
diff --git a/WORKSPACE b/WORKSPACE
index 62e88ccd52..adbc937904 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -54,14 +54,14 @@ http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/test/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcu121.zip"],
+    urls = ["https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcu121.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/test/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip"],
+    urls = ["https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip"],
 )
 
 # Download these tarballs manually from the NVIDIA website
diff --git a/docker/dist-build.sh b/docker/dist-build.sh
index ba7dd83713..02287121d8 100755
--- a/docker/dist-build.sh
+++ b/docker/dist-build.sh
@@ -3,9 +3,9 @@
 TOP_DIR=$(cd $(dirname $0); pwd)/..
 
 if [[ -z "${USE_CXX11}" ]]; then
-    BUILD_CMD="python -m pip wheel .  --extra-index-url https://download.pytorch.org/whl/test/cu121 -w dist"
+    BUILD_CMD="python -m pip wheel .  -w dist"
 else
-    BUILD_CMD="python -m pip wheel . --config-setting="--build-option=--use-cxx11-abi" --extra-index-url https://download.pytorch.org/whl/test/cu121 -w dist"
+    BUILD_CMD="python -m pip wheel . --config-setting="--build-option=--use-cxx11-abi" -w dist"
 fi
 
 # TensorRT restricts our pip version
diff --git a/py/requirements.txt b/py/requirements.txt
index b993362038..3bd06111f6 100644
--- a/py/requirements.txt
+++ b/py/requirements.txt
@@ -1,9 +1,9 @@
 numpy
 packaging
 pybind11==2.6.2
---extra-index-url https://download.pytorch.org/whl/test/cu121
+--extra-index-url https://download.pytorch.org/whl/cu121
 torch==2.2.0
-torchvision==0.16.2
+torchvision==0.17.0
 --extra-index-url https://pypi.nvidia.com
 tensorrt==8.6.1
 pyyaml
diff --git a/toolchains/ci_workspaces/WORKSPACE.x86_64.release.rhel.tmpl b/toolchains/ci_workspaces/WORKSPACE.x86_64.release.rhel.tmpl
index 058e8e946c..9d0b1073fb 100644
--- a/toolchains/ci_workspaces/WORKSPACE.x86_64.release.rhel.tmpl
+++ b/toolchains/ci_workspaces/WORKSPACE.x86_64.release.rhel.tmpl
@@ -59,14 +59,14 @@ http_archive(
     name = "libtorch",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/test/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcu121.zip"],
+    urls = ["https://download.pytorch.org/libtorch/cu121/libtorch-cxx11-abi-shared-with-deps-2.2.0%2Bcu121.zip"],
 )
 
 http_archive(
     name = "libtorch_pre_cxx11_abi",
     build_file = "@//third_party/libtorch:BUILD",
     strip_prefix = "libtorch",
-    urls = ["https://download.pytorch.org/libtorch/test/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip"],
+    urls = ["https://download.pytorch.org/libtorch/cu121/libtorch-shared-with-deps-2.2.0%2Bcu121.zip"],
 )
 
 ####################################################################################