diff --git a/nvdsinfer_custom_impl_Yolo/calibrator.h b/nvdsinfer_custom_impl_Yolo/calibrator.h
index 2a18f83..e390a76 100644
--- a/nvdsinfer_custom_impl_Yolo/calibrator.h
+++ b/nvdsinfer_custom_impl_Yolo/calibrator.h
@@ -7,6 +7,7 @@
 #define CALIBRATOR_H
 
 #include <vector>
+#include <cuda_runtime_api.h>
 
 #include "NvInfer.h"
 #include "opencv2/opencv.hpp"