From 2a9c7b5d86ac97e45c0a122fb1bfad1a8323692c Mon Sep 17 00:00:00 2001
From: xjmxyt <46990504+xjmxyt@users.noreply.github.com>
Date: Tue, 4 Apr 2023 19:32:07 +0800
Subject: [PATCH] fix set value convert out of bound (#51885)

* fix out of bound

* fix bug

* fix bug

* fix
---
 paddle/fluid/inference/tensorrt/convert/set_value_op.cc | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/inference/tensorrt/convert/set_value_op.cc b/paddle/fluid/inference/tensorrt/convert/set_value_op.cc
index b49e80c74cb..016baf012e4 100644
--- a/paddle/fluid/inference/tensorrt/convert/set_value_op.cc
+++ b/paddle/fluid/inference/tensorrt/convert/set_value_op.cc
@@ -44,7 +44,7 @@ namespace tensorrt {
 //         for h in [0,n)
 //             for w in [0,n)
 //                 output[n,c,indices[n,c,h,w],w] = updates[n,c,h,w]]
-//
+
 class SetValueConverter : public OpConverter {
  public:
   void operator()(const framework::proto::OpDesc& op,
@@ -75,6 +75,10 @@ class SetValueConverter : public OpConverter {
       axes = input_dims.nbDims - 1;
     }
 
+    if (ends == -1 || ends > input_dims.d[axes]) {
+      ends = input_dims.d[axes];
+    }
+
     if (axes >= input_dims.nbDims) {
       platform::errors::InvalidArgument(
           "The axes %d is larger than total axes %d", axes, input_dims.nbDims);
-- 
GitLab