Skip to content

Commit 2a6f999

Browse files
committed
Bug fix
Signed-off-by: Arvind Sridhar <[email protected]>
1 parent 6833017 commit 2a6f999

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

core/partitioning/partitioning.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,7 @@ std::vector<SegmentedBlock> segment_graph(torch::jit::Block* block, const Partit
319319
pytorch_nodes.clear();
320320
}
321321
if (checkLoopEvaluatable(n)) {
322-
segmented_blocks.emplace_back(SegmentedBlock::kTensorRT, std::vector<torch::jit::Node*>{n});
322+
tensorrt_nodes.push_back(n);
323323
} else {
324324
segmented_blocks.emplace_back(SegmentedBlock::kTorch, std::vector<torch::jit::Node*>{n});
325325
}

0 commit comments

Comments
 (0)