@@ -51,10 +51,11 @@ TEST(Partitioning, ExplicitNodeAutoConversionCorrectly) {
51
51
inputs_map.insert ({g->inputs ()[1 ], {inputs[1 ]}});
52
52
input_types.insert ({g->inputs ()[1 ], {{at::kInt }}});
53
53
54
- auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs (inputs_map, input_types);
55
-
54
+ partitioning_info.collection_input_spec_map = inputs_map;
56
55
torch_tensorrt::core::partitioning::PartitioningCtx ctx (g->block (), partitioning_info);
57
- torch_tensorrt::core::partitioning::partition (&ctx, input_ivalues_map);
56
+ ctx.input_types_map = input_types;
57
+ torch_tensorrt::core::partitioning::populateInputIValues (&ctx);
58
+ torch_tensorrt::core::partitioning::partition (&ctx);
58
59
auto segmented_blocks = ctx.partitioned_blocks .begin ()->second ;
59
60
60
61
for (auto & seg_block : segmented_blocks) {
@@ -93,10 +94,12 @@ TEST(Partitioning, ImplicitAutoConversionCorrectly) {
93
94
inputs_map.insert ({g->inputs ()[0 ], {inputs[0 ]}});
94
95
input_types.insert ({g->inputs ()[0 ], {{at::kFloat }}});
95
96
96
- auto input_ivalues_map = torch_tensorrt::core::partitioning::generateRandomInputs (inputs_map, input_types);
97
-
97
+ partitioning_info.collection_input_spec_map = inputs_map;
98
98
torch_tensorrt::core::partitioning::PartitioningCtx ctx (g->block (), partitioning_info);
99
- torch_tensorrt::core::partitioning::partition (&ctx, input_ivalues_map);
99
+ ctx.input_types_map = input_types;
100
+
101
+ torch_tensorrt::core::partitioning::populateInputIValues (&ctx);
102
+ torch_tensorrt::core::partitioning::partition (&ctx);
100
103
auto segmented_blocks = ctx.partitioned_blocks .begin ()->second ;
101
104
102
105
for (auto & seg_block : segmented_blocks) {
0 commit comments