diff --git a/paddle/fluid/inference/tensorrt/convert/flatten_op.cc b/paddle/fluid/inference/tensorrt/convert/flatten_op.cc
index 322b42667fa30ff31428f70c101ff81460a17df1..25351cc10ec11b733c745522499a637129d399a2 100644
--- a/paddle/fluid/inference/tensorrt/convert/flatten_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/flatten_op.cc
@@ -53,10 +53,19 @@ class FlattenOpConverter : public OpConverter {
       layer->setReshapeDimensions(flatten_dim);
     } else {
       auto* shape_layer = TRT_ENGINE_ADD_LAYER(engine_, Shape, *input);
+      nvinfer1::Dims start_dim, size_dim, stride_dim;
+      start_dim.nbDims = 1;
+      size_dim.nbDims = 1;
+      stride_dim.nbDims = 1;
+      start_dim.d[0] = 1;
+      size_dim.d[0] = dims - 1;
+      stride_dim.d[0] = 1;
+      auto* slice_layer =
+          TRT_ENGINE_ADD_LAYER(engine_, Slice, *(shape_layer->getOutput(0)),
+                               start_dim, size_dim, stride_dim);
       uint32_t reduce_dim = 1;
-
       auto* reduce_prod_layer = TRT_ENGINE_ADD_LAYER(
-          engine_, Reduce, *(shape_layer->getOutput(0)),
+          engine_, Reduce, *(slice_layer->getOutput(0)),
           nvinfer1::ReduceOperation::kPROD, reduce_dim, true);
       int32_t* constant_weight_data = new int32_t[1];
       constant_weight_data[0] = -1;
diff --git a/paddle/fluid/inference/tensorrt/helper.h b/paddle/fluid/inference/tensorrt/helper.h
index 8b557c6e6b8ce98932bc4599888d1eeec311a11e..e3c7d8b10333c322be455c1f74912b2fb11ccb75 100644
--- a/paddle/fluid/inference/tensorrt/helper.h
+++ b/paddle/fluid/inference/tensorrt/helper.h
@@ -134,6 +134,19 @@ inline size_t ProductDim(const nvinfer1::Dims& dims) {
   return v;
 }
 
+inline void PrintITensorShape(nvinfer1::ITensor* X) {
+  auto dims = X->getDimensions();
+  auto name = X->getName();
+  std::cout << "ITensor " << name << " shape: [";
+  for (int i = 0; i < dims.nbDims; i++) {
+    if (i == dims.nbDims - 1)
+      std::cout << dims.d[i];
+    else
+      std::cout << dims.d[i] << ", ";
+  }
+  std::cout << "]\n";
+}
+
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle
diff --git a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py
index 1f8f829d27c2a7a8bf7c96229cb82e2446254261..bb28fcf708503d8426f8b7f28906e7f1c444f6ad 100644
--- a/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py
+++ b/python/paddle/fluid/tests/unittests/ir/inference/test_trt_flatten_op.py
@@ -63,10 +63,10 @@ class TRTFlattenDynamicTest(InferencePassTest):
         self.trt_parameters = TRTFlattenDynamicTest.TensorRTParam(
             1 << 30, 32, 0, AnalysisConfig.Precision.Float32, False, False)
         self.dynamic_shape_params = TRTFlattenDynamicTest.DynamicShapeParam({
-            'data': [1, 6, 8, 8],
-            'flatten_0.tmp_0': [1, 6 * 8 * 8]
-        }, {'data': [3, 6, 128, 128],
-            'flatten_0.tmp_0': [3, 6 * 128 * 128]}, {
+            'data': [2, 6, 64, 64],
+            'flatten_0.tmp_0': [2, 6 * 64 * 64]
+        }, {'data': [2, 6, 64, 64],
+            'flatten_0.tmp_0': [2, 6 * 64 * 64]}, {
                 'data': [2, 6, 64, 64],
                 'flatten_0.tmp_0': [2, 6 * 64 * 64]
             }, False)