Skip to content

Commit cf58da3

Browse files
HolyWugs-olive
authored andcommitted
Install torch and tensorrt according to desired_cuda in matrix
1 parent 5a38980 commit cf58da3

File tree

2 files changed

+4
-4
lines changed

2 files changed

+4
-4
lines changed
Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
set -eou pipefail
22
source "${BUILD_ENV_FILE}"
33
# ${CONDA_RUN} ${PIP_INSTALL_TORCH} torchvision
4-
${CONDA_RUN} pip install torch torchvision --index-url https://download.pytorch.org/whl/test/cu121
4+
${CONDA_RUN} pip install torch torchvision --index-url https://download.pytorch.org/whl/test/${CU_VERSION}
55
${CONDA_RUN} pip install pyyaml mpmath==1.3.0
6-
${CONDA_RUN} pip install tensorrt==10.0.0b6 tensorrt-cu12-bindings==10.0.0b6 tensorrt-cu12-libs==10.0.0b6 --extra-index-url https://pypi.nvidia.com
6+
${CONDA_RUN} pip install tensorrt==10.0.0b6 tensorrt-${CU_VERSION::4}-bindings==10.0.0b6 tensorrt-${CU_VERSION::4}-libs==10.0.0b6 --extra-index-url https://pypi.nvidia.com
77
${CONDA_RUN} pip install ${RUNNER_ARTIFACT_DIR}/torch_tensorrt*.whl
88

99
echo -e "Running test script";

packaging/pre_build_script_windows.sh

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
python -m pip install -U numpy packaging pyyaml setuptools wheel
2-
python -m pip install torch==2.3.0 --index-url https://download.pytorch.org/whl/test/cu121
3-
python -m pip install tensorrt==10.0.0b6 tensorrt-cu12-bindings==10.0.0b6 tensorrt-cu12-libs==10.0.0b6 --extra-index-url https://pypi.nvidia.com
2+
python -m pip install torch==2.3.0 --index-url https://download.pytorch.org/whl/test/${CU_VERSION}
3+
python -m pip install tensorrt==10.0.0b6 tensorrt-${CU_VERSION::4}-bindings==10.0.0b6 tensorrt-${CU_VERSION::4}-libs==10.0.0b6 --extra-index-url https://pypi.nvidia.com
44

55
choco install bazelisk -y
66

0 commit comments

Comments
 (0)