Skip to content

Commit 18c9185

Browse files
authored
fix yolo11-seg and yolo11-pose engine cannot run after set batch > 1 (wang-xinyu#1607)
1 parent 62c1680 commit 18c9185

File tree

1 file changed

+5
-0
lines changed

1 file changed

+5
-0
lines changed

yolo11/plugin/yololayer.cu

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,11 @@ __global__ void CalDetection(const float* input, float* output, int numElements,
251251
void YoloLayerPlugin::forwardGpu(const float* const* inputs, float* output, cudaStream_t stream, int mYoloV8netHeight,
252252
int mYoloV8NetWidth, int batchSize) {
253253
int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);
254+
if (is_segmentation_) {
255+
outputElem = mMaxOutObject * (sizeof(Detection) - sizeof(float) * 51) / sizeof(float) + 1;
256+
} else if (is_pose_) {
257+
outputElem = mMaxOutObject * (sizeof(Detection) - sizeof(float) * 32) / sizeof(float) + 1;
258+
}
254259
cudaMemsetAsync(output, 0, sizeof(float), stream);
255260
for (int idx = 0; idx < batchSize; ++idx) {
256261
CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream));

0 commit comments

Comments
 (0)