From abac588e745fab66200741f45e9343b71820a311 Mon Sep 17 00:00:00 2001
From: gracehoney <31743510+aaroey@users.noreply.github.com>
Date: Tue, 6 Mar 2018 13:20:36 -0800
Subject: [PATCH] Fix the include for cuda_runtime_api.h

---
 tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc          | 2 +-
 tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc | 3 ++-
 tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h  | 9 +++++++--
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc b/tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc
index b78ff18..d4be96a 100644
--- a/tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc
+++ b/tensorflow/contrib/tensorrt/kernels/trt_calib_op.cc
@@ -25,7 +25,7 @@ limitations under the License.
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
-#include "cuda_runtime_api.h"
+#include "cuda/include/cuda_runtime_api.h"
 #include "tensorrt/include/NvInfer.h"
 
 namespace tensorflow {
diff --git a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc
index 1a842cf..1ae6347 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc
+++ b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.cc
@@ -23,7 +23,7 @@ limitations under the License.
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
-#include "cuda_runtime_api.h"
+#include "cuda/include/cuda_runtime_api.h"
 
 namespace tensorflow {
 namespace tensorrt {
@@ -120,5 +120,6 @@ TRTInt8Calibrator::~TRTInt8Calibrator() {
 
 }  // namespace tensorrt
 }  // namespace tensorflow
+
 #endif
 #endif
diff --git a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h
index aaf93ef..4e7b74d 100644
--- a/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h
+++ b/tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h
@@ -24,8 +24,10 @@ limitations under the License.
 
 #if GOOGLE_CUDA
 #if GOOGLE_TENSORRT
-#include "cuda_runtime_api.h"
+
+#include "cuda/include/cuda_runtime_api.h"
 #include "tensorrt/include/NvInfer.h"
+
 namespace tensorflow {
 namespace tensorrt {
 // This class provides a 1 element queue to match TFs push model to
@@ -61,8 +63,11 @@ struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
   bool batch_is_set_;
   string engine_name_;
 };
+
 }  // namespace tensorrt
 }  // namespace tensorflow
-#endif  // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_
+
 #endif
 #endif
+
+#endif  // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_
-- 
2.7.4