🐛 [Fix] device bug, make test can run on cpu

tests/test_utils/test_loss.py

@@ -25,11 +25,9 @@ def loss_function(cfg) -> YOLOLoss:
 
 @pytest.fixture
 def data():
-    [[torch.zeros]]
-    targets = torch.zeros(20, 6, device=torch.device("cuda"))
-    predicts = [
-        [torch.zeros(1, 144, 80 // i, 80 // i, device=torch.device("cuda")) for i in [1, 2, 4]] for _ in range(2)
-    ]
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    targets = torch.zeros(20, 6, device=device)
+    predicts = [[torch.zeros(1, 144, 80 // i, 80 // i, device=device) for i in [1, 2, 4]] for _ in range(2)]
     return predicts, targets
 
 

yolo/utils/loss.py

@@ -81,7 +81,7 @@ class YOLOLoss:
         self.class_num = cfg.hyper.data.class_num
         self.image_size = list(cfg.hyper.data.image_size)
         self.strides = cfg.model.anchor.strides
-        device = torch.device("cuda")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
         self.reverse_reg = torch.arange(self.reg_max, dtype=torch.float16, device=device)
         self.scale_up = torch.tensor(self.image_size * 2, device=device)