From 121d704827e2e3565d4b66bf7b914d15888be9a7 Mon Sep 17 00:00:00 2001
From: Robin <robin@nanovare.com>
Date: Tue, 22 Jun 2021 14:50:32 +0200
Subject: [PATCH] Fix bug described in
 https://github.com/ultralytics/yolov5/issues/2106

---
 setup.py         | 2 +-
 utils/general.py | 3 +++
 2 files changed, 4 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 0b69deacbfef..d8fd1bee9fa9 100644
--- 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
index be7c658a7b92..2c6690d4f95a 100755
--- 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