Skip to content

Commit 162f1f7

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

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -709,7 +709,8 @@ 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+
length: int = len(attrs.axes) if attrs.axes is not None else len(args[0].checked_type.shape)
713+
for i in range(0, length):
713714
begin = int(attrs.begin[i])
714715
if attrs.slice_mode == "end":
715716
end = (

0 commit comments

Comments
 (0)