diff --git a/test/ir/inference/test_save_optimized_model_pass.py b/test/ir/inference/test_save_optimized_model_pass.py
index 4e5a52285a5ee14e4acfb9793b87f039dc060778..68e2a87302b64257a46c7396845e050f72960876 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()