From c1527e4ab137e0882d3ce605432541c81d2e14a4 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 8 Jan 2020 18:48:41 -0800
Subject: [PATCH] updates

---
 test.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/test.py b/test.py
index 894ad185..4be1121d 100644
--- a/test.py
+++ b/test.py
@@ -99,7 +99,7 @@ def test(cfg,
 
             if pred is None:
                 if nl:
-                    stats.append((torch.zeros(0, 1), torch.Tensor(), torch.Tensor(), tcls))
+                    stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
                 continue
 
             # Append to text file
@@ -147,7 +147,7 @@ def test(cfg,
                             d = ti[i[j]]  # detected target
                             if d not in detected:
                                 detected.append(d)
-                                correct[pi[j]] = ious[j] > iouv  # iou_thres is 1xn
+                                correct[pi[j]] = (ious[j] > iouv).cpu()  # iou_thres is 1xn
                                 if len(detected) == nl:  # all targets already located in image
                                     break