fix: Update index in installer for CI failures (pytorch#2679)

zewenli98 · Mar 15, 2024 · 76c9ebd · 76c9ebd
1 parent ad74a73
commit 76c9ebd
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/.github/scripts/install-torch-tensorrt.sh b/.github/scripts/install-torch-tensorrt.sh
@@ -2,8 +2,9 @@
 set -eou pipefail
 # Source conda so it's available to the script environment
 source ${BUILD_ENV_FILE}
-${CONDA_RUN} ${PIP_INSTALL_TORCH} torchvision pyyaml
+${CONDA_RUN} ${PIP_INSTALL_TORCH} torchvision
+${CONDA_RUN} python -m pip install pyyaml mpmath==1.3.0
 export TRT_VERSION=$(${CONDA_RUN} python -c "import versions; versions.tensorrt_version()")
 ${CONDA_RUN} python -m pip install /opt/torch-tensorrt-builds/torch_tensorrt*+${CU_VERSION}*.whl tensorrt~=${TRT_VERSION} tensorrt-bindings~=${TRT_VERSION} --extra-index-url=https://pypi.ngc.nvidia.com
 
-echo -e "Running test script";
+echo -e "Running test script";