From a9182bdfab6c4ea0045220a5357fa1f578d4210b Mon Sep 17 00:00:00 2001
From: Abel Haro Armero <130833831+AbelHaro@users.noreply.github.com>
Date: Thu, 30 Jan 2025 10:34:26 +0100
Subject: [PATCH] Add DLA specific core usage (#18930)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
Co-authored-by: Francesco Mattioli <Francesco.mttl@gmail.com>
Co-authored-by: Lakshantha Dissanayake <lakshanthad@yahoo.com>
---
 ultralytics/nn/autobackend.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/ultralytics/nn/autobackend.py b/ultralytics/nn/autobackend.py
index f536a602..0a79b91a 100644
--- a/ultralytics/nn/autobackend.py
+++ b/ultralytics/nn/autobackend.py
@@ -293,6 +293,12 @@ class AutoBackend(nn.Module):
                 except UnicodeDecodeError:
                     f.seek(0)  # engine file may lack embedded Ultralytics metadata
                 model = runtime.deserialize_cuda_engine(f.read())  # read engine
+                if "dla" in str(device.type):
+                    dla_core = int(device.type.split(":")[1])
+                    assert dla_core in {0, 1}, (
+                        "Expected device type for inference in DLA is 'dla:0' or 'dla:1', but received '{device.type}'"
+                    )
+                    runtime.DLA_core = dla_core
 
             # Model context
             try: