From 4326f6deedbf8edfca320971ef21017a2a3ad8ac Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 5 Jul 2020 23:37:36 -0700
Subject: [PATCH] .to(device) bug fix

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

diff --git a/test.py b/test.py
index 2d6ab5454b51..1cfae9591287 100644
--- a/test.py
+++ b/test.py
@@ -31,7 +31,7 @@ def test(data,
 
         # Load model
         google_utils.attempt_download(weights)
-        model = torch.load(weights, map_location=device)['model'].float().fuse()  # load to FP32
+        model = torch.load(weights, map_location=device)['model'].float().fuse().to(device)  # load to FP32
         imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size
 
         # Multi-GPU disabled, incompatible with .half() https://github.com/ultralytics/yolov5/issues/99