From 0e4fdb6724d16034a091c5a1ad065f3801e9d351 Mon Sep 17 00:00:00 2001
From: zhoufan37 <zhoufan37@h-partners.com>
Date: Mon, 28 Feb 2022 08:57:57 +0800
Subject: [PATCH] Add YoloBugfix Operator

---
 .../test_yolo_boxes_encode.py                 | 26 ++++++++++---------
 .../aten/ops/YoloBoxesEncodeKernelNpu.cpp     |  1 +
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/test/test_network_ops/test_yolo_boxes_encode.py b/test/test_network_ops/test_yolo_boxes_encode.py
index 23656d7918..fc1330cdc2 100644
--- a/test/test_network_ops/test_yolo_boxes_encode.py
+++ b/test/test_network_ops/test_yolo_boxes_encode.py
@@ -15,8 +15,9 @@
 import torch
 import torch_npu
 
-from torch_npu.testing.testcase import TestCase, run_tests
-
+from torch_npu.testing.common_utils import TestCase, run_tests
+from torch_npu.testing.common_device_type import instantiate_device_type_tests
+from torch_npu.testing.util_test import create_common_tensor
 
 class TestYoloBoxesEncode(TestCase):
     def npu_op_exec(self, anchor_boxes, gt_bboxes, stride, impl_mode=False):
@@ -24,25 +25,26 @@ class TestYoloBoxesEncode(TestCase):
         out = out.to("cpu")
         return out.detach().numpy()
         
-    def test_yolo_boxes_encode(self, device="npu"):
-        anchor_boxes = [(2, 4)]
-        gt_bboxes = [(2 ,4)]
-        stride = [[2, 2]]
-        exoutput_list = [[[0.7921727, 0.5314963, -0.74224466, -13.815511],
+    def test_yolo_boxes_encode(self, device):
+        anchor_boxes_list = [(2, 4)]
+        gt_bboxes_list = [(2 ,4)]
+        stride_list = [[2, 2]]
+        expect_cpu_list = [[[0.7921727, 0.5314963, -0.74224466, -13.815511],
                           [0.7360072, 0.58343244, 4.3334002, -0.51378196]]]
 
-        shape_format = [[i, j, k, h] for i in anchor_boxes
-                        for j in gt_bboxes for k in stride for h in exoutput_list]
+        shape_format = [[i, j, k, h] for i in anchor_boxes_list
+                        for j in gt_bboxes_list for k in stride_list for h in expect_cpu_list]
 
         for item in shape_format:
             anchor_boxes_tensor = torch.rand(item[0], dtype=torch.float32).to("npu")
             gt_bboxes_tensor = torch.rand(item[1], dtype=torch.float32).to("npu")
-            stride_tensor = torch.tensor(item[2], dtype=torch.float32).to("npu")
-            exoutput_cpu_tensor = torch.tensor(item[3], dtype=torch.float32)
+            stride_tensor = torch.tensor(item[2], dtype=torch.int32).to("npu")
+            expect_cpu_tensor = torch.tensor(item[3], dtype=torch.float32)
             npu_output = self.npu_op_exec(anchor_boxes_tensor, gt_bboxes_tensor, stride_tensor, False)
 
-            self.assertRtolEqual(exoutput_cpu_tensor.numpy(), npu_output)
+            self.assertRtolEqual(expect_cpu_tensor.numpy(), npu_output)
 
 
+instantiate_device_type_tests(TestYoloBoxesEncode, globals(), except_for='cpu')
 if __name__ == "__main__":
     run_tests()
\ No newline at end of file
diff --git a/torch_npu/csrc/aten/ops/YoloBoxesEncodeKernelNpu.cpp b/torch_npu/csrc/aten/ops/YoloBoxesEncodeKernelNpu.cpp
index 877c0a79c0..c3263a8190 100644
--- a/torch_npu/csrc/aten/ops/YoloBoxesEncodeKernelNpu.cpp
+++ b/torch_npu/csrc/aten/ops/YoloBoxesEncodeKernelNpu.cpp
@@ -71,5 +71,6 @@ at::Tensor NPUNativeFunctions::npu_yolo_boxes_encode(
       .Run();
   return result;
 }
+
 } // namespace native
 } // namespace at_npu
\ No newline at end of file
-- 
Gitee