Add support for Torch-TensorRT in Docker

- Add install for Torch-TRT nightly
pytorch · Sep 20, 2023 · 5a1f7ec · 5a1f7ec
1 parent 8360dba
commit 5a1f7ec
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils/cuda_utils.py b/utils/cuda_utils.py
@@ -24,7 +24,7 @@
 # the numpy version needs to be consistent with
 # https://github.com/pytorch/builder/blob/e66e48f9b1968213c6a7ce3ca8df6621435f0a9c/wheel/build_wheel.sh#L146
 PIN_NUMPY_VERSION = "1.21.2"
-TORCHBENCH_TORCH_NIGHTLY_PACKAGES = ["torch", "torchvision", "torchaudio"]
+TORCHBENCH_TORCH_NIGHTLY_PACKAGES = ["torch", "torchvision", "torchaudio", "torch_tensorrt"]
 
 def _nvcc_output_match(nvcc_output, target_cuda_version):
     regex = 'release (.*),'
@@ -82,7 +82,7 @@ def install_pytorch_nightly(cuda_version: str, env, dryrun=False):
     pytorch_nightly_url = f"https://download.pytorch.org/whl/nightly/{CUDA_VERSION_MAP[cuda_version]['pytorch_url']}"
     install_torch_cmd = ["pip", "install", "--pre", "--no-cache-dir"]
     install_torch_cmd.extend(TORCHBENCH_TORCH_NIGHTLY_PACKAGES)
-    install_torch_cmd.extend(["-i",  pytorch_nightly_url])
+    install_torch_cmd.extend(["--extra-index-url",  pytorch_nightly_url])
     if dryrun:
         print(f"Install pytorch nightly: {install_torch_cmd}")
     else: