diff --git a/yolov8/yolov8_det_trt.py b/yolov8/yolov8_det_trt.py
index f65502a..f1dd1ae 100644
--- a/yolov8/yolov8_det_trt.py
+++ b/yolov8/yolov8_det_trt.py
@@ -167,7 +167,7 @@ class YoLov8TRT(object):
         # Do postprocess
         for i in range(self.batch_size):
             result_boxes, result_scores, result_classid = self.post_process(
-                output[i * 6001: (i + 1) * 6001], batch_origin_h[i], batch_origin_w[i]
+                output[i * 38001: (i + 1) * 38001], batch_origin_h[i], batch_origin_w[i]
             )
             # Draw rectangles and labels on the original image
             for j in range(len(result_boxes)):
@@ -290,7 +290,7 @@ class YoLov8TRT(object):
         # Get the num of boxes detected
         num = int(output[0])
         # Reshape to a two dimentional ndarray
-        pred = np.reshape(output[1:], (-1, 6))[:num, :]
+        pred = np.reshape(output[1:], (-1, 38))[:num, :]
         # Do nms
         boxes = self.non_max_suppression(pred, origin_h, origin_w, conf_thres=CONF_THRESH, nms_thres=IOU_THRESHOLD)
         result_boxes = boxes[:, :4] if len(boxes) else np.array([])