Update general.py

See ultralytics/yolov5#2106
alexchwong · Dec 26, 2021 · 87d2f96 · 87d2f96
1 parent 6768003
commit 87d2f96
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/utils/general.py b/utils/general.py
@@ -1089,6 +1089,9 @@ def output_to_target(output, width, height):
     targets = []
     for i, o in enumerate(output):
         if o is not None:
+            # 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()
             for pred in o:
                 box = pred[:4]
                 w = (box[2] - box[0]) / width