Fix bug described in ultralytics#2106

robin-maillot · Jun 22, 2021 · 121d704 · 121d704
1 parent 0e4f4f7
commit 121d704
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -6,7 +6,7 @@
 import shutil
 
 MODULE_NAME = "nanovare_yolov5"
-VERSION = "0.2.1"
+VERSION = "0.2.2"
 
 
 def copy_to_module():

diff --git a/utils/general.py b/utils/general.py
@@ -955,6 +955,9 @@ def output_to_target(output, width, height):
     for i, o in enumerate(output):
         if o is not None:
             for pred in o:
+                # sometimes output can be a list of tensor, so here ensure the type again, this fixes the error.
+                if isinstance(o, torch.Tensor):
+                    o = o.cpu().numpy()
                 box = pred[:4]
                 w = (box[2] - box[0]) / width
                 h = (box[3] - box[1]) / height