From 440769b954637a74d86f48eb882cc1ac4c2a94f3 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 27 Dec 2019 09:55:10 -0800
Subject: [PATCH] updates

---
 test.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/test.py b/test.py
index 642fdbcf..7cd13011 100644
--- a/test.py
+++ b/test.py
@@ -50,7 +50,9 @@ def test(cfg,
     names = load_classes(data['names'])  # class names
     # iou_thres = torch.linspace(0.5, 0.95, 10).to(device)  # for mAP@0.5:0.95
     # iou_thres = iou_thres[0].view(1)  # for mAP@0.5
-    niou = 1  # len(iou_thres)
+    if isinstance(iou_thres, float):
+        iou_thres = torch.Tensor([iou_thres]).to(device)  # convert to array
+    niou = iou_thres.numel()
 
     # Dataloader
     if dataloader is None: