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>
This commit is contained in:
parent
b91af4d3b0
commit
a9182bdfab
1 changed files with 6 additions and 0 deletions
|
|
@ -293,6 +293,12 @@ class AutoBackend(nn.Module):
|
||||||
except UnicodeDecodeError:
|
except UnicodeDecodeError:
|
||||||
f.seek(0) # engine file may lack embedded Ultralytics metadata
|
f.seek(0) # engine file may lack embedded Ultralytics metadata
|
||||||
model = runtime.deserialize_cuda_engine(f.read()) # read engine
|
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
|
# Model context
|
||||||
try:
|
try:
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue