From 9fefbed77f2f8166cacc9d4063102072f501e653 Mon Sep 17 00:00:00 2001
From: Armassarion <33727511+Armassarion@users.noreply.github.com>
Date: Wed, 22 Sep 2021 13:22:28 +0200
Subject: [PATCH] Fix a potential bug in detection number reset between each
 inference (#733)

Co-authored-by: Ksander <ksander.ejjaaouani@digeiz.com>
---
 yolov5/yololayer.cu | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/yolov5/yololayer.cu b/yolov5/yololayer.cu
index 4c042a6..0bc9f70 100644
--- 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) {