diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index ab0eeb091043d..168362e229b03 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -1505,7 +1505,8 @@ def _impl_v13(cls, inputs, attr, params): rank_input = len(infer_type(inputs[0]).checked_type.shape) num_new_axis = int(infer_type(inputs[1]).checked_type.shape[0]) - axes = relay.split(inputs[1], num_new_axis).astuple() + axes = relay.sort(inputs[1]) + axes = relay.split(axes, num_new_axis).astuple() result = inputs[0] # TODO (AndrewZhaoLuo): investigate performance issues with consecutive diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 94fd0a5de40b2..12e02d5f29c7e 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -5125,10 +5125,6 @@ def verify_eyelike(indata): "test_triu_square", "test_triu_square_neg", "test_triu_zero", - # These unsqueeze tests work, but take 2+ hrs to run - "test_unsqueeze_three_axes", - "test_unsqueeze_two_axes", - "test_unsqueeze_unsorted_axes", "test_unique_sorted_with_axis", "test_unique_sorted_with_axis_3d", "test_unique_sorted_with_negative_axis",