diff --git a/csrc/segment/normal/include/yolov8-seg.hpp b/csrc/segment/normal/include/yolov8-seg.hpp
index e6d8cad..a6a87b2 100644
--- a/csrc/segment/normal/include/yolov8-seg.hpp
+++ b/csrc/segment/normal/include/yolov8-seg.hpp
@@ -326,7 +326,7 @@ void YOLOv8_seg::postprocess(
     }
     else {
         cv::Mat matmulRes = (masks * protos).t();
-        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_w, seg_h});
+        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_h, seg_w});
 
         std::vector<cv::Mat> maskChannels;
         cv::split(maskMat, maskChannels);
diff --git a/csrc/segment/simple/include/yolov8-seg.hpp b/csrc/segment/simple/include/yolov8-seg.hpp
index d753088..25cb533 100644
--- a/csrc/segment/simple/include/yolov8-seg.hpp
+++ b/csrc/segment/simple/include/yolov8-seg.hpp
@@ -300,7 +300,7 @@ void YOLOv8_seg::postprocess(
     }
     else {
         cv::Mat matmulRes = (masks * protos).t();
-        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_w, seg_h});
+        cv::Mat maskMat   = matmulRes.reshape(indices.size(), {seg_h, seg_w});
 
         std::vector<cv::Mat> maskChannels;
         cv::split(maskMat, maskChannels);