From 431966ea3c9fe44193df9ae8b02ce93319fe28d4 Mon Sep 17 00:00:00 2001 From: ncdhz <1137436221@qq.com> Date: Thu, 1 Feb 2024 15:18:30 +0800 Subject: [PATCH] update v8 error --- src/Yolo.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/Yolo.cpp b/src/Yolo.cpp index fcb8929..dcccb4b 100644 --- a/src/Yolo.cpp +++ b/src/Yolo.cpp @@ -199,8 +199,8 @@ std::vector Yolo::prediction(torch::Tensor data) if (strcmp(this->version.c_str(), V8) == 0) { torch::Tensor pT = pred.toTensor(); - torch::Tensor score = std::get<0>(pT.slice(1, 4, -1).max(1, true)); - data = torch::cat({pT.slice(1, 0, 4), score, pT.slice(1, 4, -1)}, 1).permute({0, 2, 1}); + torch::Tensor score = std::get<0>(pT.slice(1, 4, pT.size(1)).max(1, true)); + data = torch::cat({pT.slice(1, 0, 4), score, pT.slice(1, 4, pT.size(1))}, 1).permute({0, 2, 1}); } else if (strcmp(this->version.c_str(), V6) == 0) {