From 6f5c3c8ceacd03459e050e481ae37845308b4741 Mon Sep 17 00:00:00 2001
From: Francesco Mattioli <Francesco.mttl@gmail.com>
Date: Fri, 6 Sep 2024 16:03:37 +0200
Subject: [PATCH] CoreML export: updated nms flag for end2end models (#16058)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---
 ultralytics/engine/exporter.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
index 33f245d8..f7ce9cb6 100644
--- a/ultralytics/engine/exporter.py
+++ b/ultralytics/engine/exporter.py
@@ -610,6 +610,9 @@ class Exporter:
         f = self.file.with_suffix(".mlmodel" if mlmodel else ".mlpackage")
         if f.is_dir():
             shutil.rmtree(f)
+        if self.args.nms and getattr(self.model, "end2end", False):
+            LOGGER.warning(f"{prefix} WARNING ⚠️ 'nms=True' is not available for end2end models. Forcing 'nms=False'.")
+            self.args.nms = False
 
         bias = [0.0, 0.0, 0.0]
         scale = 1 / 255