From bfc078b32fec76d425f53b0d348af255368f0f04 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Fri, 24 Feb 2023 03:22:33 -0800
Subject: [PATCH] test updates, revert Results to CPU

---
 tests/test_cli.py                  | 6 +++---
 ultralytics/yolo/engine/results.py | 6 +++---
 ultralytics/yolo/utils/plotting.py | 4 +++-
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/tests/test_cli.py b/tests/test_cli.py
index e4c1896..2879b2a 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -48,7 +48,7 @@ def test_val_classify():
 
 # Predict checks -------------------------------------------------------------------------------------------------------
 def test_predict_detect():
-    run(f"yolo predict model={MODEL}.pt source={ROOT / 'assets'} imgsz=32")
+    run(f"yolo predict model={MODEL}.pt source={ROOT / 'assets'} imgsz=32 save")
     if checks.check_online():
         run(f'yolo predict model={MODEL}.pt source=https://ultralytics.com/images/bus.jpg imgsz=32')
         run(f'yolo predict model={MODEL}.pt source=https://ultralytics.com/assets/decelera_landscape_min.mov imgsz=32')
@@ -56,11 +56,11 @@ def test_predict_detect():
 
 
 def test_predict_segment():
-    run(f"yolo predict model={MODEL}-seg.pt source={ROOT / 'assets'} imgsz=32")
+    run(f"yolo predict model={MODEL}-seg.pt source={ROOT / 'assets'} imgsz=32 save")
 
 
 def test_predict_classify():
-    run(f"yolo predict model={MODEL}-cls.pt source={ROOT / 'assets'} imgsz=32")
+    run(f"yolo predict model={MODEL}-cls.pt source={ROOT / 'assets'} imgsz=32 save")
 
 
 # Export checks --------------------------------------------------------------------------------------------------------
diff --git a/ultralytics/yolo/engine/results.py b/ultralytics/yolo/engine/results.py
index a39e5c4..12ed3f3 100644
--- a/ultralytics/yolo/engine/results.py
+++ b/ultralytics/yolo/engine/results.py
@@ -42,9 +42,9 @@ class Results:
     def __init__(self, orig_img, path, names, boxes=None, masks=None, probs=None) -> None:
         self.orig_img = orig_img
         self.orig_shape = orig_img.shape[:2]
-        self.boxes = Boxes(boxes.cpu(), self.orig_shape) if boxes is not None else None  # native size boxes
-        self.masks = Masks(masks.cpu(), self.orig_shape) if masks is not None else None  # native size or imgsz masks
-        self.probs = probs.cpu() if probs is not None else None
+        self.boxes = Boxes(boxes, self.orig_shape) if boxes is not None else None  # native size boxes
+        self.masks = Masks(masks, self.orig_shape) if masks is not None else None  # native size or imgsz masks
+        self.probs = probs if probs is not None else None
         self.names = names
         self.path = path
         self._keys = (k for k in ('boxes', 'masks', 'probs') if getattr(self, k) is not None)
diff --git a/ultralytics/yolo/utils/plotting.py b/ultralytics/yolo/utils/plotting.py
index c3b9c79..76c042c 100644
--- a/ultralytics/yolo/utils/plotting.py
+++ b/ultralytics/yolo/utils/plotting.py
@@ -114,7 +114,9 @@ class Annotator:
             self.im = np.asarray(self.im).copy()
         if len(masks) == 0:
             self.im[:] = im_gpu.permute(1, 2, 0).contiguous().cpu().numpy() * 255
-        colors = torch.tensor(colors, device=im_gpu.device, dtype=torch.float32) / 255.0
+        if im_gpu.device != masks.device:
+            im_gpu = im_gpu.to(masks.device)
+        colors = torch.tensor(colors, device=masks.device, dtype=torch.float32) / 255.0
         colors = colors[:, None, None]  # shape(n,1,1,3)
         masks = masks.unsqueeze(3)  # shape(n,h,w,1)
         masks_color = masks * (colors * alpha)  # shape(n,h,w,3)