From f502b50365292b12682dee7d1db8e83236dee957 Mon Sep 17 00:00:00 2001
From: Laughing <61612323+Laughing-q@users.noreply.github.com>
Date: Mon, 5 Jun 2023 04:35:33 +0800
Subject: [PATCH] fix overlap_mask in SegmentationLoss (#2985)

---
 ultralytics/yolo/utils/loss.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ultralytics/yolo/utils/loss.py b/ultralytics/yolo/utils/loss.py
index f1bea54..f165acb 100644
--- a/ultralytics/yolo/utils/loss.py
+++ b/ultralytics/yolo/utils/loss.py
@@ -178,10 +178,10 @@ class v8DetectionLoss:
 # Criterion class for computing training losses
 class v8SegmentationLoss(v8DetectionLoss):
 
-    def __init__(self, model, overlap=True):  # model must be de-paralleled
+    def __init__(self, model):  # model must be de-paralleled
         super().__init__(model)
         self.nm = model.model[-1].nm  # number of masks
-        self.overlap = overlap
+        self.overlap = model.args.overlap_mask
 
     def __call__(self, preds, batch):
         """Calculate and return the loss for the YOLO model."""