From 11ce877bdff2a355cc25a0c96fd5c23eb47a644b Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Wed, 8 Jan 2020 09:42:01 -0800
Subject: [PATCH] updates

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

diff --git a/utils/utils.py b/utils/utils.py
index 60efd25a..83ccd601 100755
--- a/utils/utils.py
+++ b/utils/utils.py
@@ -540,7 +540,7 @@ def non_max_suppression(prediction, conf_thres=0.5, iou_thres=0.5, multi_cls=Tru
 
         # Filter by class
         if classes:
-            pred = pred[(j.view(-1, 1) == torch.Tensor(classes)).any(1)]
+            pred = pred[(j.view(-1, 1) == torch.tensor(classes, device=j.device)).any(1)]
 
         # Apply finite constraint
         if not torch.isfinite(pred).all():