Skip to content

Commit 2c9772a

Browse files
committed
[TRT][BYOC] allow strided_slice ops on selected dimensions (#14142)
1 parent 0e046da commit 2c9772a

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

python/tvm/relay/op/contrib/tensorrt.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -709,7 +709,9 @@ def strided_slice_checker(
709709
if any([x is not None and x <= 0 for x in attrs.strides]):
710710
logger.info(f"{op_name}: stride must be positive")
711711
return False
712-
for i in range(0, len(args[0].checked_type.shape)):
712+
if len(attrs.axes) > len(args[0].checked_type.shape):
713+
logger.info(f"{op_name}: number of indexed axes is higher than input tensor rank")
714+
for i in range(0, len(attrs.axes)):
713715
begin = int(attrs.begin[i])
714716
if attrs.slice_mode == "end":
715717
end = (

0 commit comments

Comments
 (0)