Dockerfile install tensorrt-cu12==10.1.0 (ultralytics#14240)

Signed-off-by: Glenn Jocher <[email protected]>

docker/Dockerfile

@@ -35,7 +35,8 @@ ADD https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8n.pt $A
 
 # Install pip packages
 RUN python3 -m pip install --upgrade pip wheel
-RUN pip install --no-cache-dir -e ".[export]" "albumentations>=1.4.6" comet pycocotools
+# Pin TensorRT-cu12==10.1.0 to avoid 10.2.0 bug https://github.com/ultralytics/ultralytics/pull/14239 (note -cu12 must be used)
+RUN pip install --no-cache-dir -e ".[export]" "tensorrt-cu12==10.1.0" "albumentations>=1.4.6" comet pycocotools
 
 # Run exports to AutoInstall packages
 # Edge TPU export fails the first time so is run twice here

ultralytics/engine/exporter.py

@@ -688,7 +688,8 @@ def export_engine(self, prefix=colorstr("TensorRT:")):
             if LINUX:
                 check_requirements("tensorrt>7.0.0,<=10.1.0")
             import tensorrt as trt  # noqa
-        check_version(trt.__version__, "7.0.0", hard=True)  # require tensorrt>=7.0.0
+        check_version(trt.__version__, ">=7.0.0", hard=True)
+        check_version(trt.__version__, "<=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
 
         # Setup and checks
         LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")

ultralytics/nn/autobackend.py

@@ -227,7 +227,8 @@ def __init__(
                 if LINUX:
                     check_requirements("tensorrt>7.0.0,<=10.1.0")
                 import tensorrt as trt  # noqa
-            check_version(trt.__version__, "7.0.0", hard=True)  # require tensorrt>=7.0.0
+            check_version(trt.__version__, ">=7.0.0", hard=True)
+            check_version(trt.__version__, "<=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
             if device.type == "cpu":
                 device = torch.device("cuda:0")
             Binding = namedtuple("Binding", ("name", "dtype", "shape", "data", "ptr"))