From 1646b8de4b8f7e4dcea54e8683c668d853cca7ca Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 4 Aug 2019 01:07:54 +0200
Subject: [PATCH] updates

---
 test.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test.py b/test.py
index 63d89f3b..9dbb738e 100644
--- a/test.py
+++ b/test.py
@@ -73,7 +73,7 @@ def test(cfg,
 
         # Compute loss
         if hasattr(model, 'hyp'):  # if model has loss hyperparameters
-            loss += compute_loss(train_out, targets, model)[1][[0, 2, 3]]  # GIoU, obj, cls
+            loss += compute_loss(train_out, targets, model)[1][[0, 2, 3]].cpu()  # GIoU, obj, cls
 
         # Run NMS
         output = non_max_suppression(inf_out, conf_thres=conf_thres, nms_thres=nms_thres)