diff --git a/python/tvm/relay/frontend/tflite.py b/python/tvm/relay/frontend/tflite.py index 1b593ad8dea3..d6f704703cae 100644 --- a/python/tvm/relay/frontend/tflite.py +++ b/python/tvm/relay/frontend/tflite.py @@ -3093,7 +3093,7 @@ def convert_detection_postprocess(self, op): valid_count = ret[0] # keep only the top 'max_detections' rows ret = _op.strided_slice( - ret[1], [0, 0, 0], [batch_size, custom_options["max_detections"], anchor_boxes] + ret[1], [0, 0, 0], [batch_size, custom_options["max_detections"], 6] ) # the output needs some reshaping to match tflite ret = _op.split(ret, 6, axis=2)