experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Merge pull request #25771 from trevor-m:tmorris_tftrt_calibrator_trt_5_1

Browse Source 
PiperOrigin-RevId: 234747254
This commit is contained in:
TensorFlower Gardener 2019-02-20 00:42:03 -08:00
commit 8501ea4c41
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
tensorflow/compiler/tf2tensorrt/utils/trt_int8_calibrator.h
View File

@ -34,7 +34,12 @@ namespace tensorrt {
// TRTs pull model for calibration. When TRT implements a means for // TRTs pull model for calibration. When TRT implements a means for
// a push calibration This class should be updated accordingly // 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 { struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
#endif
public: public:
// Construct a calibrator for future calibration. // Construct a calibrator for future calibration.
TRTInt8Calibrator( TRTInt8Calibrator(