From 29efbdb6b5d182b97bab404309298162093f7008 Mon Sep 17 00:00:00 2001 From: Yuanle Liu Date: Tue, 22 Aug 2023 17:02:07 +0800 Subject: [PATCH] Fix error of test_save_optimized_model_pass unittest (#56522) --- test/ir/inference/test_save_optimized_model_pass.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/test/ir/inference/test_save_optimized_model_pass.py b/test/ir/inference/test_save_optimized_model_pass.py index 4e5a52285a5..68e2a87302b 100644 --- a/test/ir/inference/test_save_optimized_model_pass.py +++ b/test/ir/inference/test_save_optimized_model_pass.py @@ -122,9 +122,9 @@ class TestSaveOptimizedModelPassWithTRT( use_calib_mode=False, ) config.set_trt_dynamic_shape_info( - {"x": [1, 3, 224, 224], "flatten_1.tmp_0": [1, 9216]}, - {"x": [1, 3, 224, 224], "flatten_1.tmp_0": [1, 9216]}, - {"x": [1, 3, 224, 224], "flatten_1.tmp_0": [1, 9216]}, + {"x": [1, 3, 224, 224], "flatten_0.tmp_0": [1, 9216]}, + {"x": [1, 3, 224, 224], "flatten_0.tmp_0": [1, 9216]}, + {"x": [1, 3, 224, 224], "flatten_0.tmp_0": [1, 9216]}, ) config.exp_disable_tensorrt_ops(["flatten_contiguous_range"]) config.enable_memory_optim() -- GitLab