From 3c4a3207222c987bafd0e583c4cb0b621368e884 Mon Sep 17 00:00:00 2001 From: Trevor Morris Date: Thu, 14 Feb 2019 16:33:56 -0800 Subject: [PATCH] Use IInt8EntropyCalibrator2 for TRT 5.1 onward --- tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h index 10587e99624..aa70b07f8d7 100644 --- a/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h +++ b/tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h @@ -34,7 +34,12 @@ namespace tensorrt { // TRTs pull model for calibration. When TRT implements a means for // a push calibration This class should be updated accordingly +// IInt8EntropyCalibrator2 is prefferred for TRT 5.1+. +#if NV_TENSORRT_MAJOR > 5 || (NV_TENSORRT_MAJOR == 5 && NV_TENSORRT_MINOR >= 1) +struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator2 { +#else struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator { +#endif public: // Construct a calibrator for future calibration. TRTInt8Calibrator(