From c2ecaa7573b3cecade5503c64362597d13c4e101 Mon Sep 17 00:00:00 2001 From: Bo Wang Date: Thu, 17 Nov 2022 22:34:59 -0800 Subject: [PATCH] fix: fix failed test cases caused by partition API changes Signed-off-by: Bo Wang --- tests/core/partitioning/test_shape_analysis.cpp | 2 +- .../partitioning/test_type_auto_conversion.cpp | 15 +++++++++------ 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/core/partitioning/test_shape_analysis.cpp b/tests/core/partitioning/test_shape_analysis.cpp index 7558fbec69..a1cb70876d 100644 --- a/tests/core/partitioning/test_shape_analysis.cpp +++ b/tests/core/partitioning/test_shape_analysis.cpp @@ -11,7 +11,7 @@ bool checkSegmentedBlockInputShape( if (segmented_blocks.size() != in_shape.size()) return false; for (size_t i = 0; i < segmented_blocks.size(); ++i) { - auto cur_block_in_shapes = segmented_blocks[i].in_shapes(); + auto cur_block_in_shapes = segmented_blocks[i].in_opt_shapes(); if (cur_block_in_shapes.size() != in_shape[i].size()) return false; for (size_t j = 0; j < cur_block_in_shapes.size(); ++j) { diff --git a/tests/core/partitioning/test_type_auto_conversion.cpp b/tests/core/partitioning/test_type_auto_conversion.cpp index d7b7e2391f..28f620b843 100644 --- a/tests/core/partitioning/test_type_auto_conversion.cpp +++ b/tests/core/partitioning/test_type_auto_conversion.cpp @@ -51,10 +51,11 @@ TEST(Partitioning, ExplicitNodeAutoConversionCorrectly) { inputs_map.insert({g->inputs()[1], {inputs[1]}}); input_types.insert({g->inputs()[1], {{at::kInt}}}); - auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types); - + partitioning_info.collection_input_spec_map = inputs_map; torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info); - torch_tensorrt::core::partitioning::partition(&ctx, input_ivalues_map); + ctx.input_types_map = input_types; + torch_tensorrt::core::partitioning::populateInputIValues(&ctx); + torch_tensorrt::core::partitioning::partition(&ctx); auto segmented_blocks = ctx.partitioned_blocks.begin()->second; for (auto& seg_block : segmented_blocks) { @@ -93,10 +94,12 @@ TEST(Partitioning, ImplicitAutoConversionCorrectly) { inputs_map.insert({g->inputs()[0], {inputs[0]}}); input_types.insert({g->inputs()[0], {{at::kFloat}}}); - auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs(inputs_map, input_types); - + partitioning_info.collection_input_spec_map = inputs_map; torch_tensorrt::core::partitioning::PartitioningCtx ctx(g->block(), partitioning_info); - torch_tensorrt::core::partitioning::partition(&ctx, input_ivalues_map); + ctx.input_types_map = input_types; + + torch_tensorrt::core::partitioning::populateInputIValues(&ctx); + torch_tensorrt::core::partitioning::partition(&ctx); auto segmented_blocks = ctx.partitioned_blocks.begin()->second; for (auto& seg_block : segmented_blocks) {