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

Fix Multi-GPU trained model export (#14551)

Browse source 
Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Mohammed Yasin 2024-07-20 03:01:11 +08:00 • committed by GitHub
parent dcde8bd23d
commit 0d059bec0c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 7 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

8
ultralytics/engine/model.py
View file

@ -588,7 +588,13 @@ class Model(nn.Module):
self._check_is_pytorch_model() self._check_is_pytorch_model()
from .exporter import Exporter from .exporter import Exporter
custom = {"imgsz": self.model.args["imgsz"], "batch": 1, "data": None, "verbose": False} # method defaults custom = {
"imgsz": self.model.args["imgsz"],
"batch": 1,
"data": None,
"device": None, # reset to avoid multi-GPU errors
"verbose": False,
} # method defaults
args = {**self.overrides, **custom, **kwargs, "mode": "export"} # highest priority args on the right args = {**self.overrides, **custom, **kwargs, "mode": "export"} # highest priority args on the right
return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model) return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model)