ultralytics 8.3.199__tar.gz → 8.3.200__tar.gz

{ultralytics-8.3.199 → ultralytics-8.3.200}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ultralytics
-Version: 8.3.199
+Version: 8.3.200
 Summary: Ultralytics YOLO 🚀 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author-email: Glenn Jocher <glenn.jocher@ultralytics.com>, Jing Qiu <jing.qiu@ultralytics.com>
 Maintainer-email: Ultralytics <hello@ultralytics.com>

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
-__version__ = "8.3.199"
+__version__ = "8.3.200"
 
 import importlib
 import os

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/engine/exporter.py

@@ -106,7 +106,7 @@ from ultralytics.utils.checks import (
     is_sudo_available,
 )
 from ultralytics.utils.downloads import attempt_download_asset, get_github_assets, safe_download
-from ultralytics.utils.export import export_engine, export_onnx
+from ultralytics.utils.export import onnx2engine, torch2imx, torch2onnx
 from ultralytics.utils.files import file_size, spaces_in_path
 from ultralytics.utils.metrics import batch_probiou
 from ultralytics.utils.nms import TorchNMS
@@ -284,7 +284,8 @@ class Exporter:
             # Get the closest match if format is invalid
             matches = difflib.get_close_matches(fmt, fmts, n=1, cutoff=0.6)  # 60% similarity required to match
             if not matches:
-                raise ValueError(f"Invalid export format='{fmt}'. Valid formats are {fmts}")
+                msg = "Model is already in PyTorch format." if fmt == "pt" else f"Invalid export format='{fmt}'."
+                raise ValueError(f"{msg} Valid formats are {fmts}")
             LOGGER.warning(f"Invalid export format='{fmt}', updating to format='{matches[0]}'")
             fmt = matches[0]
         flags = [x == fmt for x in fmts]
@@ -408,9 +409,9 @@ class Exporter:
         model = model.fuse()
 
         if imx:
-            from ultralytics.utils.torch_utils import FXModel
+            from ultralytics.utils.export.imx import FXModel
 
-            model = FXModel(model)
+            model = FXModel(model, self.imgsz)
         for m in model.modules():
             if isinstance(m, Classify):
                 m.export = True
@@ -425,15 +426,6 @@ class Exporter:
             elif isinstance(m, C2f) and not is_tf_format:
                 # EdgeTPU does not support FlexSplitV while split provides cleaner ONNX graph
                 m.forward = m.forward_split
-            if isinstance(m, Detect) and imx:
-                from ultralytics.utils.tal import make_anchors
-
-                m.anchors, m.strides = (
-                    x.transpose(0, 1)
-                    for x in make_anchors(
-                        torch.cat([s / m.stride.unsqueeze(-1) for s in self.imgsz], dim=1), m.stride, 0.5
-                    )
-                )
 
         y = None
         for _ in range(2):  # dry runs
@@ -609,7 +601,7 @@ class Exporter:
             self.args.opset = opset_version  # for NMSModel
 
         with arange_patch(self.args):
-            export_onnx(
+            torch2onnx(
                 NMSModel(self.model, self.args) if self.args.nms else self.model,
                 self.im,
                 f,
@@ -932,7 +924,7 @@ class Exporter:
         LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...")
         assert Path(f_onnx).exists(), f"failed to export ONNX file: {f_onnx}"
         f = self.file.with_suffix(".engine")  # TensorRT engine file
-        export_engine(
+        onnx2engine(
             f_onnx,
             f,
             self.args.workspace,
@@ -1168,7 +1160,6 @@ class Exporter:
     @try_export
     def export_imx(self, prefix=colorstr("IMX:")):
         """Export YOLO model to IMX format."""
-        gptq = False
         assert LINUX, (
             "export only supported on Linux. "
             "See https://developer.aitrios.sony-semicon.com/en/raspberrypi-ai-camera/documentation/imx500-converter"
@@ -1181,13 +1172,6 @@ class Exporter:
         check_requirements("imx500-converter[pt]>=3.16.1")  # Separate requirements for imx500-converter
         check_requirements("mct-quantizers>=1.6.0")  # Separate for compatibility with model-compression-toolkit
 
-        import model_compression_toolkit as mct
-        import onnx
-        from edgemdt_tpc import get_target_platform_capabilities
-        from sony_custom_layers.pytorch import multiclass_nms_with_indices
-
-        LOGGER.info(f"\n{prefix} starting export with model_compression_toolkit {mct.__version__}...")
-
         # Install Java>=17
         try:
             java_output = subprocess.run(["java", "--version"], check=True, capture_output=True).stdout.decode()
@@ -1198,150 +1182,17 @@ class Exporter:
             cmd = (["sudo"] if is_sudo_available() else []) + ["apt", "install", "-y", "openjdk-21-jre"]
             subprocess.run(cmd, check=True)
 
-        def representative_dataset_gen(dataloader=self.get_int8_calibration_dataloader(prefix)):
-            for batch in dataloader:
-                img = batch["img"]
-                img = img / 255.0
-                yield [img]
-
-        tpc = get_target_platform_capabilities(tpc_version="4.0", device_type="imx500")
-
-        bit_cfg = mct.core.BitWidthConfig()
-        if "C2PSA" in self.model.__str__():  # YOLO11
-            if self.model.task == "detect":
-                layer_names = ["sub", "mul_2", "add_14", "cat_21"]
-                weights_memory = 2585350.2439
-                n_layers = 238  # 238 layers for fused YOLO11n
-            elif self.model.task == "pose":
-                layer_names = ["sub", "mul_2", "add_14", "cat_22", "cat_23", "mul_4", "add_15"]
-                weights_memory = 2437771.67
-                n_layers = 257  # 257 layers for fused YOLO11n-pose
-        else:  # YOLOv8
-            if self.model.task == "detect":
-                layer_names = ["sub", "mul", "add_6", "cat_17"]
-                weights_memory = 2550540.8
-                n_layers = 168  # 168 layers for fused YOLOv8n
-            elif self.model.task == "pose":
-                layer_names = ["add_7", "mul_2", "cat_19", "mul", "sub", "add_6", "cat_18"]
-                weights_memory = 2482451.85
-                n_layers = 187  # 187 layers for fused YOLO11n-pose
-
-        # Check if the model has the expected number of layers
-        if len(list(self.model.modules())) != n_layers:
-            raise ValueError("IMX export only supported for YOLOv8n and YOLO11n models.")
-
-        for layer_name in layer_names:
-            bit_cfg.set_manual_activation_bit_width([mct.core.common.network_editors.NodeNameFilter(layer_name)], 16)
-
-        config = mct.core.CoreConfig(
-            mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=10),
-            quantization_config=mct.core.QuantizationConfig(concat_threshold_update=True),
-            bit_width_config=bit_cfg,
-        )
-
-        resource_utilization = mct.core.ResourceUtilization(weights_memory=weights_memory)
-
-        quant_model = (
-            mct.gptq.pytorch_gradient_post_training_quantization(  # Perform Gradient-Based Post Training Quantization
-                model=self.model,
-                representative_data_gen=representative_dataset_gen,
-                target_resource_utilization=resource_utilization,
-                gptq_config=mct.gptq.get_pytorch_gptq_config(
-                    n_epochs=1000, use_hessian_based_weights=False, use_hessian_sample_attention=False
-                ),
-                core_config=config,
-                target_platform_capabilities=tpc,
-            )[0]
-            if gptq
-            else mct.ptq.pytorch_post_training_quantization(  # Perform post training quantization
-                in_module=self.model,
-                representative_data_gen=representative_dataset_gen,
-                target_resource_utilization=resource_utilization,
-                core_config=config,
-                target_platform_capabilities=tpc,
-            )[0]
-        )
-
-        class NMSWrapper(torch.nn.Module):
-            """Wrap PyTorch Module with multiclass_nms layer from sony_custom_layers."""
-
-            def __init__(
-                self,
-                model: torch.nn.Module,
-                score_threshold: float = 0.001,
-                iou_threshold: float = 0.7,
-                max_detections: int = 300,
-                task: str = "detect",
-            ):
-                """
-                Initialize NMSWrapper with PyTorch Module and NMS parameters.
-
-                Args:
-                    model (torch.nn.Module): Model instance.
-                    score_threshold (float): Score threshold for non-maximum suppression.
-                    iou_threshold (float): Intersection over union threshold for non-maximum suppression.
-                    max_detections (int): The number of detections to return.
-                    task (str): Task type, either 'detect' or 'pose'.
-                """
-                super().__init__()
-                self.model = model
-                self.score_threshold = score_threshold
-                self.iou_threshold = iou_threshold
-                self.max_detections = max_detections
-                self.task = task
-
-            def forward(self, images):
-                """Forward pass with model inference and NMS post-processing."""
-                # model inference
-                outputs = self.model(images)
-
-                boxes, scores = outputs[0], outputs[1]
-                nms_outputs = multiclass_nms_with_indices(
-                    boxes=boxes,
-                    scores=scores,
-                    score_threshold=self.score_threshold,
-                    iou_threshold=self.iou_threshold,
-                    max_detections=self.max_detections,
-                )
-                if self.task == "pose":
-                    kpts = outputs[2]  # (bs, max_detections, kpts 17*3)
-                    out_kpts = torch.gather(kpts, 1, nms_outputs.indices.unsqueeze(-1).expand(-1, -1, kpts.size(-1)))
-                    return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, out_kpts
-                return nms_outputs.boxes, nms_outputs.scores, nms_outputs.labels, nms_outputs.n_valid
-
-        quant_model = NMSWrapper(
-            model=quant_model,
-            score_threshold=self.args.conf or 0.001,
-            iou_threshold=self.args.iou,
-            max_detections=self.args.max_det,
-            task=self.model.task,
-        ).to(self.device)
-
-        f = Path(str(self.file).replace(self.file.suffix, "_imx_model"))
-        f.mkdir(exist_ok=True)
-        onnx_model = f / Path(str(self.file.name).replace(self.file.suffix, "_imx.onnx"))  # js dir
-        mct.exporter.pytorch_export_model(
-            model=quant_model, save_model_path=onnx_model, repr_dataset=representative_dataset_gen
-        )
-
-        model_onnx = onnx.load(onnx_model)  # load onnx model
-        for k, v in self.metadata.items():
-            meta = model_onnx.metadata_props.add()
-            meta.key, meta.value = k, str(v)
-
-        onnx.save(model_onnx, onnx_model)
-
-        subprocess.run(
-            ["imxconv-pt", "-i", str(onnx_model), "-o", str(f), "--no-input-persistency", "--overwrite-output"],
-            check=True,
+        return torch2imx(
+            self.model,
+            self.file,
+            self.args.conf,
+            self.args.iou,
+            self.args.max_det,
+            metadata=self.metadata,
+            dataset=self.get_int8_calibration_dataloader(prefix),
+            prefix=prefix,
         )
 
-        # Needed for imx models.
-        with open(f / "labels.txt", "w", encoding="utf-8") as file:
-            file.writelines([f"{name}\n" for _, name in self.model.names.items()])
-
-        return f
-
     def _add_tflite_metadata(self, file):
         """Add metadata to *.tflite models per https://ai.google.dev/edge/litert/models/metadata."""
         import zipfile
@@ -1531,7 +1382,7 @@ class NMSModel(torch.nn.Module):
         for i in range(bs):
             box, cls, score, extra = boxes[i], classes[i], scores[i], extras[i]
             mask = score > self.args.conf
-            if self.is_tf:
+            if self.is_tf or (self.args.format == "onnx" and self.obb):
                 # TFLite GatherND error if mask is empty
                 score *= mask
                 # Explicit length otherwise reshape error, hardcoded to `self.args.max_det * 5`

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/engine/trainer.py

@@ -174,7 +174,22 @@ class BaseTrainer:
 
         # Callbacks
         self.callbacks = _callbacks or callbacks.get_default_callbacks()
-        if RANK in {-1, 0}:
+
+        if isinstance(self.args.device, str) and len(self.args.device):  # i.e. device='0' or device='0,1,2,3'
+            world_size = len(self.args.device.split(","))
+        elif isinstance(self.args.device, (tuple, list)):  # i.e. device=[0, 1, 2, 3] (multi-GPU from CLI is list)
+            world_size = len(self.args.device)
+        elif self.args.device in {"cpu", "mps"}:  # i.e. device='cpu' or 'mps'
+            world_size = 0
+        elif torch.cuda.is_available():  # i.e. device=None or device='' or device=number
+            world_size = 1  # default to device 0
+        else:  # i.e. device=None or device=''
+            world_size = 0
+
+        self.ddp = world_size > 1 and "LOCAL_RANK" not in os.environ
+        self.world_size = world_size
+        # Run subprocess if DDP training, else train normally
+        if RANK in {-1, 0} and not self.ddp:
             callbacks.add_integration_callbacks(self)
             # Start console logging immediately at trainer initialization
             self.run_callbacks("on_pretrain_routine_start")
@@ -194,19 +209,8 @@ class BaseTrainer:
 
     def train(self):
         """Allow device='', device=None on Multi-GPU systems to default to device=0."""
-        if isinstance(self.args.device, str) and len(self.args.device):  # i.e. device='0' or device='0,1,2,3'
-            world_size = len(self.args.device.split(","))
-        elif isinstance(self.args.device, (tuple, list)):  # i.e. device=[0, 1, 2, 3] (multi-GPU from CLI is list)
-            world_size = len(self.args.device)
-        elif self.args.device in {"cpu", "mps"}:  # i.e. device='cpu' or 'mps'
-            world_size = 0
-        elif torch.cuda.is_available():  # i.e. device=None or device='' or device=number
-            world_size = 1  # default to device 0
-        else:  # i.e. device=None or device=''
-            world_size = 0
-
         # Run subprocess if DDP training, else train normally
-        if world_size > 1 and "LOCAL_RANK" not in os.environ:
+        if self.ddp:
             # Argument checks
             if self.args.rect:
                 LOGGER.warning("'rect=True' is incompatible with Multi-GPU training, setting 'rect=False'")
@@ -218,7 +222,7 @@ class BaseTrainer:
                 self.args.batch = 16
 
             # Command
-            cmd, file = generate_ddp_command(world_size, self)
+            cmd, file = generate_ddp_command(self)
             try:
                 LOGGER.info(f"{colorstr('DDP:')} debug command {' '.join(cmd)}")
                 subprocess.run(cmd, check=True)
@@ -228,7 +232,7 @@ class BaseTrainer:
                 ddp_cleanup(self, str(file))
 
         else:
-            self._do_train(world_size)
+            self._do_train()
 
     def _setup_scheduler(self):
         """Initialize training learning rate scheduler."""
@@ -238,20 +242,19 @@ class BaseTrainer:
             self.lf = lambda x: max(1 - x / self.epochs, 0) * (1.0 - self.args.lrf) + self.args.lrf  # linear
         self.scheduler = optim.lr_scheduler.LambdaLR(self.optimizer, lr_lambda=self.lf)
 
-    def _setup_ddp(self, world_size):
+    def _setup_ddp(self):
         """Initialize and set the DistributedDataParallel parameters for training."""
         torch.cuda.set_device(RANK)
         self.device = torch.device("cuda", RANK)
-        # LOGGER.info(f'DDP info: RANK {RANK}, WORLD_SIZE {world_size}, DEVICE {self.device}')
         os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1"  # set to enforce timeout
         dist.init_process_group(
             backend="nccl" if dist.is_nccl_available() else "gloo",
             timeout=timedelta(seconds=10800),  # 3 hours
             rank=RANK,
-            world_size=world_size,
+            world_size=self.world_size,
         )
 
-    def _setup_train(self, world_size):
+    def _setup_train(self):
         """Build dataloaders and optimizer on correct rank process."""
         ckpt = self.setup_model()
         self.model = self.model.to(self.device)
@@ -293,13 +296,13 @@ class BaseTrainer:
             callbacks_backup = callbacks.default_callbacks.copy()  # backup callbacks as check_amp() resets them
             self.amp = torch.tensor(check_amp(self.model), device=self.device)
             callbacks.default_callbacks = callbacks_backup  # restore callbacks
-        if RANK > -1 and world_size > 1:  # DDP
+        if RANK > -1 and self.world_size > 1:  # DDP
             dist.broadcast(self.amp.int(), src=0)  # broadcast from rank 0 to all other ranks; gloo errors with boolean
         self.amp = bool(self.amp)  # as boolean
         self.scaler = (
             torch.amp.GradScaler("cuda", enabled=self.amp) if TORCH_2_4 else torch.cuda.amp.GradScaler(enabled=self.amp)
         )
-        if world_size > 1:
+        if self.world_size > 1:
             self.model = nn.parallel.DistributedDataParallel(self.model, device_ids=[RANK], find_unused_parameters=True)
 
         # Check imgsz
@@ -312,7 +315,7 @@ class BaseTrainer:
             self.args.batch = self.batch_size = self.auto_batch()
 
         # Dataloaders
-        batch_size = self.batch_size // max(world_size, 1)
+        batch_size = self.batch_size // max(self.world_size, 1)
         self.train_loader = self.get_dataloader(
             self.data["train"], batch_size=batch_size, rank=LOCAL_RANK, mode="train"
         )
@@ -350,11 +353,11 @@ class BaseTrainer:
         self.scheduler.last_epoch = self.start_epoch - 1  # do not move
         self.run_callbacks("on_pretrain_routine_end")
 
-    def _do_train(self, world_size=1):
+    def _do_train(self):
         """Train the model with the specified world size."""
-        if world_size > 1:
-            self._setup_ddp(world_size)
-        self._setup_train(world_size)
+        if self.world_size > 1:
+            self._setup_ddp()
+        self._setup_train()
 
         nb = len(self.train_loader)  # number of batches
         nw = max(round(self.args.warmup_epochs * nb), 100) if self.args.warmup_epochs > 0 else -1  # warmup iterations
@@ -365,7 +368,7 @@ class BaseTrainer:
         self.run_callbacks("on_train_start")
         LOGGER.info(
             f"Image sizes {self.args.imgsz} train, {self.args.imgsz} val\n"
-            f"Using {self.train_loader.num_workers * (world_size or 1)} dataloader workers\n"
+            f"Using {self.train_loader.num_workers * (self.world_size or 1)} dataloader workers\n"
             f"Logging results to {colorstr('bold', self.save_dir)}\n"
             f"Starting training for " + (f"{self.args.time} hours..." if self.args.time else f"{self.epochs} epochs...")
         )
@@ -417,7 +420,7 @@ class BaseTrainer:
                     loss, self.loss_items = unwrap_model(self.model).loss(batch, preds)
                     self.loss = loss.sum()
                     if RANK != -1:
-                        self.loss *= world_size
+                        self.loss *= self.world_size
                     self.tloss = (
                         (self.tloss * i + self.loss_items) / (i + 1) if self.tloss is not None else self.loss_items
                     )

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/engine/tuner.py

@@ -341,7 +341,8 @@ class Tuner:
             hyp[k] = round(min(max(hyp[k], bounds[0]), bounds[1]), 5)
 
         # Update types
-        hyp["close_mosaic"] = int(round(hyp["close_mosaic"]))
+        if "close_mosaic" in hyp:
+            hyp["close_mosaic"] = int(round(hyp["close_mosaic"]))
 
         return hyp
 

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/nn/modules/head.py

@@ -162,7 +162,7 @@ class Detect(nn.Module):
         # Inference path
         shape = x[0].shape  # BCHW
         x_cat = torch.cat([xi.view(shape[0], self.no, -1) for xi in x], 2)
-        if self.format != "imx" and (self.dynamic or self.shape != shape):
+        if self.dynamic or self.shape != shape:
             self.anchors, self.strides = (x.transpose(0, 1) for x in make_anchors(x, self.stride, 0.5))
             self.shape = shape
 
@@ -182,8 +182,6 @@ class Detect(nn.Module):
             dbox = self.decode_bboxes(self.dfl(box) * norm, self.anchors.unsqueeze(0) * norm[:, :2])
         else:
             dbox = self.decode_bboxes(self.dfl(box), self.anchors.unsqueeze(0)) * self.strides
-        if self.export and self.format == "imx":
-            return dbox.transpose(1, 2), cls.sigmoid().permute(0, 2, 1)
         return torch.cat((dbox, cls.sigmoid()), 1)
 
     def bias_init(self):
@@ -387,8 +385,6 @@ class Pose(Detect):
         if self.training:
             return x, kpt
         pred_kpt = self.kpts_decode(bs, kpt)
-        if self.export and self.format == "imx":
-            return (*x, pred_kpt.permute(0, 2, 1))
         return torch.cat([x, pred_kpt], 1) if self.export else (torch.cat([x[0], pred_kpt], 1), (x[1], kpt))
 
     def kpts_decode(self, bs: int, kpts: torch.Tensor) -> torch.Tensor:

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/utils/benchmarks.py

@@ -145,7 +145,7 @@ def benchmark(
                 assert not is_end2end
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 IMX exports not supported"
                 assert model.task == "detect", "IMX only supported for detection task"
-                assert "C2f" in model.__str__(), "IMX only supported for YOLOv8"  # TODO: enable for YOLO11
+                assert "C2f" in model.__str__(), "IMX only supported for YOLOv8n and YOLO11n"
             if format == "rknn":
                 assert not isinstance(model, YOLOWorld), "YOLOWorldv2 RKNN exports not supported yet"
                 assert not is_end2end, "End-to-end models not supported by RKNN yet"

{ultralytics-8.3.199 → ultralytics-8.3.200}/ultralytics/utils/dist.py

@@ -76,12 +76,11 @@ if __name__ == "__main__":
     return file.name
 
 
-def generate_ddp_command(world_size: int, trainer):
+def generate_ddp_command(trainer):
     """
     Generate command for distributed training.
 
     Args:
-        world_size (int): Number of processes to spawn for distributed training.
         trainer (ultralytics.engine.trainer.BaseTrainer): The trainer containing configuration for distributed training.
 
     Returns:
@@ -95,7 +94,16 @@ def generate_ddp_command(world_size: int, trainer):
     file = generate_ddp_file(trainer)
     dist_cmd = "torch.distributed.run" if TORCH_1_9 else "torch.distributed.launch"
     port = find_free_network_port()
-    cmd = [sys.executable, "-m", dist_cmd, "--nproc_per_node", f"{world_size}", "--master_port", f"{port}", file]
+    cmd = [
+        sys.executable,
+        "-m",
+        dist_cmd,
+        "--nproc_per_node",
+        f"{trainer.world_size}",
+        "--master_port",
+        f"{port}",
+        file,
+    ]
     return cmd, file
 
 

ultralytics-8.3.199/ultralytics/utils/export.py → ultralytics-8.3.200/ultralytics/utils/export/__init__.py

@@ -9,8 +9,10 @@ import torch
 
 from ultralytics.utils import IS_JETSON, LOGGER
 
+from .imx import torch2imx  # noqa
 
-def export_onnx(
+
+def torch2onnx(
     torch_model: torch.nn.Module,
     im: torch.Tensor,
     onnx_file: str,
@@ -47,7 +49,7 @@ def export_onnx(
     )
 
 
-def export_engine(
+def onnx2engine(
     onnx_file: str,
     engine_file: str | None = None,
     workspace: int | None = None,
@@ -98,12 +100,12 @@ def export_engine(
     # Engine builder
     builder = trt.Builder(logger)
     config = builder.create_builder_config()
-    workspace = int((workspace or 0) * (1 << 30))
+    workspace_bytes = int((workspace or 0) * (1 << 30))
     is_trt10 = int(trt.__version__.split(".", 1)[0]) >= 10  # is TensorRT >= 10
-    if is_trt10 and workspace > 0:
-        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace)
-    elif workspace > 0:  # TensorRT versions 7, 8
-        config.max_workspace_size = workspace
+    if is_trt10 and workspace_bytes > 0:
+        config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace_bytes)
+    elif workspace_bytes > 0:  # TensorRT versions 7, 8
+        config.max_workspace_size = workspace_bytes
     flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
     network = builder.create_network(flag)
     half = builder.platform_has_fast_fp16 and half