Merge branch 'training'

Author: lus105

src/data/classification_datamodule.py

@@ -123,21 +123,21 @@ def setup(self, stage: Optional[str] = None) -> None:
 
             if hasattr(self.data_train, 'classes') and self._class_names is None:
                 self._class_names = self.data_train.classes
-            
+
             self.setup_stages_done.add('fit')
 
         elif stage == 'predict':
             if 'predict' in self.setup_stages_done:
                 return
-            
+
             self.data_predict = ImageFolder(
                 root=Path(self.test_data_dir),
                 transform=self.val_test_transforms,
             )
 
             if hasattr(self.data_predict, 'classes') and self._class_names is None:
                 self._class_names = self.data_predict.classes
-            
+
             self.setup_stages_done.add('predict')
 
     def train_dataloader(self) -> DataLoader[Any]:

src/models/components/base_model.py

@@ -63,10 +63,10 @@ def __init__(
         self.model_repo = model_repo
         self.model_kwargs = kwargs
         self.model = None
-    
+
     def load_model(self, **additional_kwargs: Any) -> None:
         """Load or reload the model with saved kwargs and optional additional kwargs.
-        
+
         Args:
             **additional_kwargs: Additional keyword arguments to override or extend saved kwargs.
         """

src/models/mnist_module.py

@@ -48,6 +48,7 @@ def __init__(
         scheduler: torch.optim.lr_scheduler,
         compile: bool,
         ckpt_path: str,
+        num_classes: int = 10,
     ) -> None:
         """Initialize a `MNISTLitModule`.
 

src/train.py

@@ -48,7 +48,9 @@ def train(cfg: DictConfig) -> tuple[dict[str, Any], dict[str, Any]]:
     datamodule: LightningDataModule = hydra.utils.instantiate(cfg.data)
 
     log.info(f'Instantiating model <{cfg.model._target_}>')
-    model: LightningModule = hydra.utils.instantiate(cfg.model, num_classes=datamodule.num_classes)
+    model: LightningModule = hydra.utils.instantiate(
+        cfg.model, num_classes=datamodule.num_classes
+    )
 
     log.info('Instantiating loggers...')
     loggers: list[Logger] = instantiate_loggers(cfg.get('logger'))