From 0bf40070afe832e137328241a0ef8a0b5c4b7768 Mon Sep 17 00:00:00 2001
From: zhoutianzi666 <39978853+zhoutianzi666@users.noreply.github.com>
Date: Thu, 25 Aug 2022 10:07:15 +0800
Subject: [PATCH] enforce_reshape (#45386)

---
 paddle/fluid/inference/tensorrt/convert/reshape_op.cc | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/paddle/fluid/inference/tensorrt/convert/reshape_op.cc b/paddle/fluid/inference/tensorrt/convert/reshape_op.cc
index 5ad9bdcae7..5ad2f9eefc 100644
--- a/paddle/fluid/inference/tensorrt/convert/reshape_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/reshape_op.cc
@@ -70,6 +70,16 @@ class ReshapeOpConverter : public OpConverter {
       layer->setReshapeDimensions(reshape_dim);
     else
       layer->setInput(1, *real_shape_tensor);
+
+    PADDLE_ENFORCE_GE(
+        layer->getOutput(0)->getDimensions().nbDims,
+        0,
+        platform::errors::InvalidArgument(
+            "Errors occures in Paddle-TRT reshape2 op, try to use C++ Api "
+            "config.Exp_DisableTensorRtOPs({\"reshape2\"})\n; or Python Api "
+            "config.exp_disable_tensorrt_ops([\"reshape2\"]) to forbid "
+            "reshape2 op into "
+            "Paddle-TRT."));
     auto output_name = op_desc.Output("Out")[0];
     RreplenishLayerAndOutput(layer, "reshape", {output_name}, test_mode);
   }
-- 
GitLab