🐛 [Fix] the default value of function input

tests/test_model/test_yolo.py

@@ -20,7 +20,7 @@ def test_build_model():
 
         OmegaConf.set_struct(cfg.model, False)
         cfg.weight = None
-        model = YOLO(cfg.model, 80)
+        model = YOLO(cfg.model)
         assert len(model.model) == 38
 
 

yolo/tools/data_augmentation.py

@@ -7,7 +7,7 @@ from torchvision.transforms import functional as TF
 class AugmentationComposer:
     """Composes several transforms together."""
 
-    def __init__(self, transforms, image_size: int = 640):
+    def __init__(self, transforms, image_size: int = [640, 640]):
         self.transforms = transforms
         self.image_size = image_size[0]
         self.pad_resize = PadAndResize(self.image_size)