HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Fix DLA export when device=None (#17128)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
This commit is contained in:
Laughing 2024-10-24 17:58:10 +08:00 • committed by GitHub
parent 8f0a94409f
commit be40a45ec3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

7
ultralytics/engine/exporter.py
View file

@ -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