From 18c9185591c71d9654f77f13d48b1efb7ad9ec78 Mon Sep 17 00:00:00 2001
From: mpj1234 <78143172+mpj1234@users.noreply.github.com>
Date: Fri, 6 Dec 2024 11:50:22 +0800
Subject: [PATCH] fix yolo11-seg and yolo11-pose engine cannot run after set
 batch > 1 (#1607)

---
 yolo11/plugin/yololayer.cu | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/yolo11/plugin/yololayer.cu b/yolo11/plugin/yololayer.cu
index c3ae07f..54cdb89 100644
--- a/yolo11/plugin/yololayer.cu
+++ b/yolo11/plugin/yololayer.cu
@@ -251,6 +251,11 @@ __global__ void CalDetection(const float* input, float* output, int numElements,
 void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,
                                  int mYoloV8NetWidth, int batchSize) {
     int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);
+    if (is_segmentation_) {
+        outputElem = mMaxOutObject * (sizeof(Detection) - sizeof(float) * 51) / sizeof(float) + 1;
+    } else if (is_pose_) {
+        outputElem = mMaxOutObject * (sizeof(Detection) - sizeof(float) * 32) / sizeof(float) + 1;
+    }
     cudaMemsetAsync(output, 0, sizeof(float), stream);
     for (int idx = 0; idx < batchSize; ++idx) {
         CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream));