fix(//cpp/ptq): Enable FP16 kernels for INT8 applications

Signed-off-by: Naren Dasan <naren@narendasan.com> Signed-off-by: Naren Dasan <narens@nvidia.com>
pytorch · Jun 11, 2020 · 26709cc · 26709cc
1 parent aa131ac
commit 26709cc
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/core/conversion/conversionctx/ConversionCtx.cpp b/core/conversion/conversionctx/ConversionCtx.cpp
@@ -51,6 +51,7 @@ ConversionCtx::ConversionCtx(BuilderSettings build_settings)
     case nvinfer1::DataType::kINT8:
         TRTORCH_CHECK(builder->platformHasFastInt8(), "Requested inference in INT8 but platform does support INT8");
         cfg->setFlag(nvinfer1::BuilderFlag::kINT8);
+        cfg->setFlag(nvinfer1::BuilderFlag::kFP16);
         input_type = nvinfer1::DataType::kFLOAT;
         TRTORCH_CHECK(settings.calibrator != nullptr, "Requested inference in INT8 but no calibrator provided, set the ptq_calibrator field in the ExtraInfo struct with your calibrator");
         cfg->setInt8Calibrator(settings.calibrator);