Separating the function to loop over data from the one to create batches for training

src/config/configuration.py

@@ -106,6 +106,7 @@ class TrainCfg:
 class DataCfg:
     name: str
     path: str
+    batch_size: int = 0
 
 
 @dataclass(frozen=True)

src/driver/continuous_monitor.py

@@ -132,7 +132,7 @@ def _process_stream(self) -> None:
         Raises:
             StopIteration: When the data loader is exhausted
         """
-        train_loader, val_loader = self.modelHarness.get_cur_data_loaders()
+        train_loader, val_loader = self.modelHarness.get_cur_loop_loaders()
 
         for batch_idx, batch in tqdm(
             enumerate(val_loader),

src/model/torch_model_harness.py

@@ -55,6 +55,7 @@ def update_data_stream(self) -> None:
     def get_cur_data_loaders(self) -> Tuple[DataLoader, DataLoader]:
         """
         Returns a training and validation dataloader compatible with the model input
+        that will be used for continual learning
         """
         raise NotImplementedError
 
@@ -88,6 +89,14 @@ def _unpack(self, batch: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
         x, y = batch
         return x, y
 
+    @torch.no_grad()
+    def get_cur_loop_loaders(self) -> Tuple[DataLoader, DataLoader]:
+        """
+        Returns a training and validation dataloader compatible with the model input
+        that will be used to loop over for inference
+        """
+        return self.get_cur_data_loaders()
+
     @staticmethod
     def _to_scalar(x: Tensor | float) -> float:
         if isinstance(x, torch.Tensor):