🔀 [Merge] branch 'TEST' into DOCKER

tests/test_tools/test_data_loader.py

@@ -66,4 +66,4 @@ def test_directory_stream_data_loader_frame(directory_stream_data_loader: Stream
     frame, rev_tensor, origin_frame = next(iter(directory_stream_data_loader))
     assert frame.shape == (1, 3, 640, 640)
     assert rev_tensor.shape == (1, 5)
-    assert origin_frame.size == (480, 640)
+    assert origin_frame.size == (480, 640) or origin_frame.size == (512, 640)

tests/test_tools/test_dataset_preparation.py

@@ -22,7 +22,8 @@ def test_prepare_dataset(train_cfg: Config):
         assert len(os.listdir(data_type)) == 5
 
     annotations_path = Path("tests/data/annotations")
-    assert os.listdir(annotations_path) == ["instances_val.json", "instances_train.json"]
+    assert "instances_val.json" in os.listdir(annotations_path)
+    assert "instances_train.json" in os.listdir(annotations_path)
 
 
 def test_prepare_weight():

yolo/utils/model_utils.py

@@ -94,7 +94,8 @@ def get_device(device_spec: Union[str, int, List[int]]) -> torch.device:
         device_spec = initialize_distributed()
     if torch.cuda.is_available() and "cuda" in str(device_spec):
         return torch.device(device_spec), ddp_flag
-    return torch.device("cpu"), False
+    device = torch.device(device_spec)
+    return device, ddp_flag
 
 
 class PostProccess: