[Paddle-TRT]TensorRT version check in index_put op (#67720)

ckl117 · web-flow · commit 54fee119cec2 · 2024-08-28T20:12:46.000+08:00
diff --git a/paddle/fluid/inference/tensorrt/convert/index_put_op.cc b/paddle/fluid/inference/tensorrt/convert/index_put_op.cc
@@ -21,7 +21,7 @@ class IndexPutOpConverter : public OpConverter {
   void operator()(const framework::proto::OpDesc& op,
                   const framework::Scope& scope,
                   bool test_mode) override {
-#if IS_TRT_VERSION_GE(8200)
+#if IS_TRT_VERSION_GE(8510)
     VLOG(3) << "convert a index_put op to tensorrt layer";
     framework::OpDesc op_desc(op, nullptr);
     std::string input_name = op_desc.Input("x").front();
@@ -67,11 +67,8 @@ class IndexPutOpConverter : public OpConverter {
     indices_slice_layer->setInput(1, *start_tensor);
     indices_slice_layer->setInput(2, *input_shape_tensor);
     indices_slice_layer->setInput(3, *stride_tensor);
-#if IS_TRT_VERSION_GE(8500)
     indices_slice_layer->setMode(nvinfer1::SampleMode::kCLAMP);
-#else
-    indices_slice_layer->setMode(nvinfer1::SliceMode::kCLAMP);
-#endif
+
     auto* bool_indices_tensor =
         Cast(indices_slice_layer->getOutput(0), nvinfer1::DataType::kBOOL);
     // nonzero
@@ -100,11 +97,8 @@ class IndexPutOpConverter : public OpConverter {
     value_slice_layer->setInput(1, *Add1DConstantLayer(0));
     value_slice_layer->setInput(2, *indices_count_tensor);
     value_slice_layer->setInput(3, *Add1DConstantLayer(1));
-#if IS_TRT_VERSION_GE(8500)
     value_slice_layer->setMode(nvinfer1::SampleMode::kCLAMP);
-#else
-    value_slice_layer->setMode(nvinfer1::SliceMode::kCLAMP);
-#endif
+
     value_tensor = value_slice_layer->getOutput(0);
     auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
                                        Scatter,
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -2857,8 +2857,8 @@ struct SimpleOpTypeSetTeller : public Teller {
     }
 
     if (op_type == "index_put") {
-#if IS_TRT_VERSION_LT(8200)
-      VLOG(3) << "index_put is not supported when TensorRT < 8.2";
+#if IS_TRT_VERSION_LT(8510)
+      VLOG(3) << "index_put is not supported when TensorRT < 8.5.1";
       return false;
 #endif
       if (!with_dynamic_shape) {
