fix: erase the repetitive nodes in dependency analysis

Signed-off-by: Bo Wang <[email protected]>
pytorch · Apr 27, 2021 · 80b1038 · 80b1038
1 parent 4a318a2
commit 80b1038
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/core/partitioning/partitioning.cpp b/core/partitioning/partitioning.cpp
@@ -97,12 +97,14 @@ std::vector<SegmentedBlock> injectNodesForNonTensorInputs(SegmentedBlock& seg_bl
       if (containNonTensorInputs(n, nontensor_inputs_set) || prev_non_tensor_outputs) {
         if (!tensorrt_nodes.empty()) {
           new_seg_blocks.emplace_back(SegmentedBlock::kTensorRT, tensorrt_nodes);
+          tensorrt_nodes.clear();
         }
         pytorch_nodes.push_back(n);
         prev_non_tensor_outputs = containNonTensorOutputs(n);
       } else {
         if (!pytorch_nodes.empty()) {
           new_seg_blocks.emplace_back(SegmentedBlock::kTorch, pytorch_nodes);
+          pytorch_nodes.clear();
         }
         tensorrt_nodes.push_back(n);
       }