Skip to content

Commit

Permalink
Fix typo in YOLOv8-Libtorch-CPP-Inference (#9330)
Browse files Browse the repository at this point in the history
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
  • Loading branch information
higor-melo and glenn-jocher committed Mar 27, 2024
1 parent 898cbcd commit 1325889
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions examples/YOLOv8-LibTorch-CPP-Inference/main.cc
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ torch::Tensor nms(const torch::Tensor& bboxes, const torch::Tensor& scores, floa
}


torch::Tensor non_max_supperession(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45, int max_det = 300) {
torch::Tensor non_max_suppression(torch::Tensor& prediction, float conf_thres = 0.25, float iou_thres = 0.45, int max_det = 300) {
auto bs = prediction.size(0);
auto nc = prediction.size(1) - 4;
auto nm = prediction.size(1) - nc - 4;
Expand Down Expand Up @@ -237,7 +237,7 @@ int main() {
torch::Tensor output = yolo_model.forward(inputs).toTensor().cpu();

// NMS
auto keep = non_max_supperession(output)[0];
auto keep = non_max_suppression(output)[0];
auto boxes = keep.index({Slice(), Slice(None, 4)});
keep.index_put_({Slice(), Slice(None, 4)}, scale_boxes({input_image.rows, input_image.cols}, boxes, {image.rows, image.cols}));

Expand Down

0 comments on commit 1325889

Please sign in to comment.