Fix a potential bug in detection number reset between each inference (wang-xinyu#733)

Armassarion · KsanderE · web-flow · commit 9fefbed77f2f · 2021-09-22T19:22:28.000+08:00
Co-authored-by: Ksander &lt;ksander.ejjaaouani@digeiz.com&gt;
diff --git a/yolov5/yololayer.cu b/yolov5/yololayer.cu
@@ -237,7 +237,7 @@ namespace nvinfer1
     {
         int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);
         for (int idx = 0; idx < batchSize; ++idx) {
-            CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float)));
+            CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream));
         }
         int numElem = 0;
         for (unsigned int i = 0; i < mYoloKernel.size(); ++i) {

Original file line number	Diff line number	Diff line change
`@@ -237,7 +237,7 @@ namespace nvinfer1`
`237`	`237`	`{`
`238`	`238`	`int outputElem = 1 + mMaxOutObject * sizeof(Detection) / sizeof(float);`
`239`	`239`	`for (int idx = 0; idx < batchSize; ++idx) {`
`240`		`- CUDA_CHECK(cudaMemset(output + idx * outputElem, 0, sizeof(float)));`
	`240`	`+ CUDA_CHECK(cudaMemsetAsync(output + idx * outputElem, 0, sizeof(float), stream));`
`241`	`241`	`}`
`242`	`242`	`int numElem = 0;`
`243`	`243`	`for (unsigned int i = 0; i < mYoloKernel.size(); ++i) {`