From abbc33925cfc132ce1a1400cf1df179c8f62c049 Mon Sep 17 00:00:00 2001
From: zhoutianzi666 <39978853+zhoutianzi666@users.noreply.github.com>
Date: Mon, 29 May 2023 15:36:47 +0800
Subject: [PATCH] fix cumsum bug (#54160)

---
 paddle/fluid/inference/tensorrt/convert/cumsum_op.cc   | 2 +-
 paddle/fluid/inference/tensorrt/convert/op_converter.h | 6 ++++++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/cumsum_op.cc b/paddle/fluid/inference/tensorrt/convert/cumsum_op.cc
index 3f5c186d1ab..d8777d5edca 100644
--- a/paddle/fluid/inference/tensorrt/convert/cumsum_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/cumsum_op.cc
@@ -51,7 +51,7 @@ class CumsumOpConverter : public OpConverter {
             return Add1DConstantLayer(d, "", scalar);
           } else {
             nvinfer1::ITensor* inpShape = Shape(inpTensor);
-            return GetEleTensorOfShape(inpShape, d, scalar);
+            return GetEleTensorOfShape(inpShape, axis, scalar);
           }
         };
 
diff --git a/paddle/fluid/inference/tensorrt/convert/op_converter.h b/paddle/fluid/inference/tensorrt/convert/op_converter.h
index d61495bbb3f..a7c569338b4 100644
--- a/paddle/fluid/inference/tensorrt/convert/op_converter.h
+++ b/paddle/fluid/inference/tensorrt/convert/op_converter.h
@@ -529,6 +529,12 @@ class OpConverter {
   nvinfer1::ITensor* GetEleTensorOfShape(nvinfer1::ITensor* shape_tensor,
                                          int index,
                                          bool is_scalar = false) {
+    PADDLE_ENFORCE_GE(
+        index,
+        0,
+        platform::errors::PreconditionNotMet(
+            "The index should be greater or equal than 0, but got %d", index));
+
     auto* tensor =
         TRT_ENGINE_ADD_LAYER(engine_,
                              Gather,
-- 
GitLab