diff --git a/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/dataset.py b/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/dataset.py
index c050f32..d67cd5f 100644
--- a/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/dataset.py
+++ b/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/dataset.py
@@ -136,7 +136,9 @@ def make_grid(anchors, nx=20, ny=20):
 
 def correct_bbox(result, anchors, stride, cls_num, out):
     result = torch.tensor(result)
-    bs, _, ny, nx, _ = result.shape
+    bs, _, ny, nx = result.shape
+    result = result.reshape(-1, 3, 85, ny, nx)
+    result = result.permute(0, 1, 3, 4, 2)
     grid, anchor_grid = make_grid(anchors, nx, ny)
     y = result.float().sigmoid()
     y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + grid) * stride  # xy
diff --git a/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/model.py b/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/model.py
index ec6f2c6..5215518 100644
--- a/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/model.py
+++ b/built-in/ACL_Pytorch/Yolov5_for_Pytorch/common/util/model.py
@@ -20,7 +20,7 @@ from pathlib import Path
 from common.util.dataset import coco80_to_coco91_class, correct_bbox, save_coco_json
 
 try:
-    from utils.general import non_max_suppression, scale_coords  # tag > 2.0
+    from utils.general import non_max_suppression, scale_boxes  # tag > 2.0
 except:
     from utils.utils import non_max_suppression, scale_coords  # tag = 2.0
 
@@ -83,7 +83,7 @@ def forward_nms_script(model, dataloader, cfg):
         boxout = nms(box_out, conf_thres=cfg["conf_thres"], iou_thres=cfg["iou_thres"])
         for idx, pred in enumerate(boxout):
             try:
-                scale_coords(img[idx].shape[1:], pred[:, :4], shapes[idx][0], shapes[idx][1])  # native-space pred
+                scale_boxes(img[idx].shape[1:], pred[:, :4], shapes[idx][0], shapes[idx][1])  # native-space pred
             except:
                 pred = torch.tensor([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0]])
             # append to COCO-JSON dictionary