Fix Multi-GPU trained model export (#14551)
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
parent
dcde8bd23d
commit
0d059bec0c
1 changed files with 7 additions and 1 deletions
|
|
@ -588,7 +588,13 @@ class Model(nn.Module):
|
|||
self._check_is_pytorch_model()
|
||||
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
|
||||
return Exporter(overrides=args, _callbacks=self.callbacks)(model=self.model)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue