From 0c815aef8035aefe2a2b7eb0102c942e054f2d54 Mon Sep 17 00:00:00 2001 From: Marcos Luciano Date: Fri, 18 Feb 2022 03:51:08 -0300 Subject: [PATCH] Update yoloPlugins.cpp --- nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp b/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp index 8350567..69de790 100644 --- a/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp +++ b/nvdsinfer_custom_impl_Yolo/yoloPlugins.cpp @@ -212,8 +212,8 @@ int YoloLayer::enqueue( } else { void* softmax; - cudaMallocHost(&softmax, sizeof(outputs[0])); - cudaMemcpy(softmax, outputs[0], sizeof(outputs[0]), cudaMemcpyHostToDevice); + CHECK(cudaMallocHost(&softmax, sizeof(outputs[0]))); + CHECK(cudaMemcpy(softmax, outputs[0], sizeof(outputs[0]), cudaMemcpyHostToDevice)); CHECK(cudaYoloLayer_v2( inputs[0], outputs[0], softmax, batchSize, m_GridSizeX, m_GridSizeY, m_NumClasses, m_NumBoxes, @@ -276,4 +276,4 @@ nvinfer1::IPluginV2* YoloLayer::clone() const noexcept return new YoloLayer (m_NumBoxes, m_NumClasses, m_GridSizeX, m_GridSizeY, m_Type, m_NewCoords, m_ScaleXY, m_BetaNMS, m_Anchors, m_Mask); } -REGISTER_TENSORRT_PLUGIN(YoloLayerPluginCreator); \ No newline at end of file +REGISTER_TENSORRT_PLUGIN(YoloLayerPluginCreator);