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);