diff --git a/docker/Dockerfile b/docker/Dockerfile index e59825eabc..be157f7967 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -4,7 +4,7 @@ # Start FROM PyTorch image https://hub.docker.com/r/pytorch/pytorch or nvcr.io/nvidia/pytorch:23.03-py3 FROM pytorch/pytorch:2.2.2-cuda12.1-cudnn8-runtime -RUN pip install --no-cache-dir nvidia-tensorrt --index-url https://pypi.ngc.nvidia.com +RUN pip install --no-cache-dir tensorrt # Set environment variables ENV APP_HOME /usr/src/ultralytics diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 5f8953f1b2..f47f67cec3 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -686,7 +686,7 @@ class Exporter: import tensorrt as trt # noqa except ImportError: if LINUX: - check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com") + check_requirements("tensorrt", cmds="-U") import tensorrt as trt # noqa check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0 diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py index 70acd28c40..216b126013 100644 --- a/ultralytics/nn/autobackend.py +++ b/ultralytics/nn/autobackend.py @@ -225,7 +225,7 @@ class AutoBackend(nn.Module): import tensorrt as trt # noqa https://developer.nvidia.com/nvidia-tensorrt-download except ImportError: if LINUX: - check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com") + check_requirements("tensorrt", cmds="-U") import tensorrt as trt # noqa check_version(trt.__version__, "7.0.0", hard=True) # require tensorrt>=7.0.0 if device.type == "cpu":