diff --git a/tests/python/unittest/test_meta_schedule_integration.py b/tests/python/unittest/test_meta_schedule_integration.py index ad0d4832732f7..1c3ef8ae0e192 100644 --- a/tests/python/unittest/test_meta_schedule_integration.py +++ b/tests/python/unittest/test_meta_schedule_integration.py @@ -69,24 +69,24 @@ def test_meta_schedule_integration_extract_from_resnet(): expected_task_names = [ "fused_" + s for s in [ - "nn_max_pool2d", - "nn_adaptive_avg_pool2d", - "nn_dense_add", - "nn_conv2d_add", - "nn_conv2d_add_1", "nn_conv2d_add_2", - "nn_conv2d_add_add_nn_relu", + "nn_conv2d_add_1", + "nn_conv2d_add", + "nn_conv2d_add_nn_relu_7", + "nn_max_pool2d", + "nn_conv2d_add_nn_relu_6", + "nn_conv2d_add_add_nn_relu_3", + "nn_conv2d_add_nn_relu_5", + "nn_conv2d_add_nn_relu_4", + "nn_conv2d_add_add_nn_relu_2", + "nn_conv2d_add_nn_relu_3", + "nn_conv2d_add_nn_relu_2", "nn_conv2d_add_add_nn_relu_1", - "nn_conv2d_add_nn_relu", "nn_conv2d_add_nn_relu_1", - "nn_conv2d_add_nn_relu_2", - "nn_conv2d_add_nn_relu_3", - "nn_conv2d_add_nn_relu_4", - "nn_conv2d_add_nn_relu_5", - "nn_contrib_conv2d_winograd_without_weight_transform_add_add_nn_relu", - "nn_contrib_conv2d_winograd_without_weight_transform_add_add_nn_relu_1", - "nn_contrib_conv2d_winograd_without_weight_transform_add_nn_relu", - "nn_contrib_conv2d_winograd_without_weight_transform_add_nn_relu_1", + "nn_conv2d_add_nn_relu", + "nn_conv2d_add_add_nn_relu", + "nn_adaptive_avg_pool2d", + "nn_contrib_dense_pack_add", # The two tasks below are purely spatial and are ruled out by AutoScheduler "layout_transform", "layout_transform_reshape_squeeze", @@ -95,8 +95,7 @@ def test_meta_schedule_integration_extract_from_resnet(): assert len(extracted_tasks) == 20 for t in extracted_tasks: - print(t.task_name) - # assert t.task_name in expected_task_names, t.task_name + assert t.task_name in expected_task_names, t.task_name @requires_torch @@ -152,6 +151,4 @@ def print_results(self) -> None: if __name__ == "__main__": - # sys.exit(pytest.main([__file__] + sys.argv[1:])) - # test_meta_schedule_integration_extract_from_resnet() - test_meta_schedule_integration_apply_history_best() + sys.exit(pytest.main([__file__] + sys.argv[1:]))