config/args added

config/args.py

@@ -0,0 +1,23 @@
+from dataclasses import dataclass
+
+
+@dataclass
+class Args:
+    """
+    Training arguments.
+    """
+
+    # Learning rate for the optimizer
+    learning_rate: float = 1e-3
+    # Training batch size
+    batch_size: int = 64
+    # Total numebr of classes
+    num_classes: int = 10
+    # Maximum number of training epochs
+    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