diff --git a/nvdsinfer_custom_impl_Yolo/yoloPlugins.h b/nvdsinfer_custom_impl_Yolo/yoloPlugins.h
index 7d41791..1293796 100644
--- a/nvdsinfer_custom_impl_Yolo/yoloPlugins.h
+++ b/nvdsinfer_custom_impl_Yolo/yoloPlugins.h
@@ -26,6 +26,8 @@
 #ifndef __YOLO_PLUGINS__
 #define __YOLO_PLUGINS__
 
+#include <cuda_runtime_api.h>
+
 #include "yolo.h"
 
 #define CUDA_CHECK(status) {                                                                                               \