From 81d7c740d2758a55df23f5ed879696f311b69af6 Mon Sep 17 00:00:00 2001
From: gs-olive <113141689+gs-olive@users.noreply.github.com>
Date: Fri, 26 Apr 2024 11:18:47 -0700
Subject: [PATCH] Add CUDA-specific TRT build

---
 .github/scripts/install-torch-tensorrt-windows.sh | 2 +-
 packaging/pre_build_script_windows.sh             | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/.github/scripts/install-torch-tensorrt-windows.sh b/.github/scripts/install-torch-tensorrt-windows.sh
index 58d6647111..97a5f53b16 100644
--- a/.github/scripts/install-torch-tensorrt-windows.sh
+++ b/.github/scripts/install-torch-tensorrt-windows.sh
@@ -7,7 +7,7 @@ ${CONDA_RUN} pip install pyyaml mpmath==1.3.0
 
 # Install TRT from PyPi
 TRT_VERSION=$(${CONDA_RUN} python -c "import yaml; print(yaml.safe_load(open('dev_dep_versions.yml', 'r'))['__tensorrt_version__'])")
-${CONDA_RUN} pip install tensorrt==${TRT_VERSION} tensorrt-${CU_VERSION::4}-bindings==${TRT_VERSION} tensorrt-${CU_VERSION::4}-libs==${TRT_VERSION} --extra-index-url https://pypi.nvidia.com
+${CONDA_RUN} pip install tensorrt==${TRT_VERSION} tensorrt-${CU_VERSION::4}==${TRT_VERSION} tensorrt-${CU_VERSION::4}-bindings==${TRT_VERSION} tensorrt-${CU_VERSION::4}-libs==${TRT_VERSION} --extra-index-url https://pypi.nvidia.com
 
 # Install pre-built Torch-TRT
 ${CONDA_RUN} pip install ${RUNNER_ARTIFACT_DIR}/torch_tensorrt*.whl
diff --git a/packaging/pre_build_script_windows.sh b/packaging/pre_build_script_windows.sh
index 10d1c1a56b..5f49488e56 100644
--- a/packaging/pre_build_script_windows.sh
+++ b/packaging/pre_build_script_windows.sh
@@ -3,7 +3,7 @@ python -m pip install -U numpy packaging pyyaml setuptools wheel
 # Install TRT from PyPi
 TRT_VERSION=$(${CONDA_RUN} python -c "import yaml; print(yaml.safe_load(open('dev_dep_versions.yml', 'r'))['__tensorrt_version__'])")
 
-python -m pip install tensorrt==${TRT_VERSION} tensorrt-${CU_VERSION::4}-bindings==${TRT_VERSION} tensorrt-${CU_VERSION::4}-libs==${TRT_VERSION} --extra-index-url https://pypi.nvidia.com
+python -m pip install tensorrt==${TRT_VERSION} tensorrt-${CU_VERSION::4}==${TRT_VERSION} tensorrt-${CU_VERSION::4}-bindings==${TRT_VERSION} tensorrt-${CU_VERSION::4}-libs==${TRT_VERSION} --extra-index-url https://pypi.nvidia.com
 
 choco install bazelisk -y