diff --git a/models/yolov5s.yaml b/models/yolov5s.yaml
index f35beabb1e1c..a5729d247ee9 100644
--- a/models/yolov5s.yaml
+++ b/models/yolov5s.yaml
@@ -44,5 +44,5 @@ head:
    [[-1, 10], 1, Concat, [1]],  # cat head P5
    [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)
 
-   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
+  [[17, 20, 23, 2], 1, SemanticSegment, [nc, anchors]],  # Detect(P3, P4, P5)
   ]