From 0d059bec0c49a3c9a3333b41ad1cb0431b5ee399 Mon Sep 17 00:00:00 2001 From: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com> Date: Sat, 20 Jul 2024 03:01:11 +0800 Subject: [PATCH] Fix Multi-GPU trained model export (#14551) Co-authored-by: UltralyticsAssistant Co-authored-by: Glenn Jocher --- ultralytics/engine/model.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/ultralytics/engine/model.py b/ultralytics/engine/model.py index f8bc495f..e626f8ca 100644 --- a/ultralytics/engine/model.py +++ b/ultralytics/engine/model.py @@ -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)