From 3d74666e7ffa12781bd72e3198fb7f19789df7bc Mon Sep 17 00:00:00 2001 From: dotsimulate <95192856+dotsimulate@users.noreply.github.com> Date: Wed, 7 Feb 2024 17:54:30 -0500 Subject: [PATCH] Update install-tensorrt.py updates TensorRT installation to select tensort post version based on the CUDA version --- src/streamdiffusion/tools/install-tensorrt.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/streamdiffusion/tools/install-tensorrt.py b/src/streamdiffusion/tools/install-tensorrt.py index 182871c4..b2a32c3d 100644 --- a/src/streamdiffusion/tools/install-tensorrt.py +++ b/src/streamdiffusion/tools/install-tensorrt.py @@ -31,7 +31,7 @@ def install(cu: Optional[Literal["11", "12"]] = get_cuda_version_from_torch()): if not is_installed("tensorrt"): run_pip(f"install {cudnn_name} --no-cache-dir") run_pip( - "install --pre --extra-index-url https://pypi.nvidia.com tensorrt==9.0.1.post11.dev4 --no-cache-dir" + f"install --pre --extra-index-url https://pypi.nvidia.com tensorrt==9.0.1.post{cu}.dev4 --no-cache-dir" ) if not is_installed("polygraphy"):