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

View file

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