args fixed

config/args.py

@@ -10,8 +10,14 @@ class Args:
     # Learning rate for the optimizer
     learning_rate: float = 1e-3
     # Training batch size
-    batch_size: int = 32
+    batch_size: int = 64
     # Total numebr of classes
     num_classes: int = 10
     # Maximum number of training epochs
-    max_epochs: int = 5
+    max_epochs: int = 100
+    # Input shape
+    input_shape: tuple = (3, 224, 224)
+    # Use pretrained weights
+    # Can be "IMAGENET1K_V1", "IMAGENET1K_V2", "DEFAULT"
+    # CHec more at https://pytorch.org/vision/stable/models.html
+    weights: str = None