From be40a45ec32d9173ff42b4bdb00dbfe5a8c7a838 Mon Sep 17 00:00:00 2001 From: Laughing <61612323+Laughing-q@users.noreply.github.com> Date: Thu, 24 Oct 2024 17:58:10 +0800 Subject: [PATCH] Fix DLA export when device=None (#17128) Co-authored-by: UltralyticsAssistant --- ultralytics/engine/exporter.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index da2e746cbe..6d403a2afb 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -195,13 +195,12 @@ class Exporter: # Device dla = None - if fmt == "engine" and "dla" in self.args.device: - dla = self.args.device.split(":")[-1] - assert dla in {"0", "1"}, f"Expected self.args.device='dla:0' or 'dla:1, but got {self.args.device}." - self.args.device = "0" if fmt == "engine" and self.args.device is None: LOGGER.warning("WARNING ⚠️ TensorRT requires GPU export, automatically assigning device=0") self.args.device = "0" + if fmt == "engine" and "dla" in str(self.args.device): # convert int/list to str first + dla = self.args.device.split(":")[-1] + assert dla in {"0", "1"}, f"Expected self.args.device='dla:0' or 'dla:1, but got {self.args.device}." self.device = select_device("cpu" if self.args.device is None else self.args.device) # Checks