More modularization

1 files changed, 9 insertions, 6 deletions

diff --git a/training/util.py b/training/util.py
index 0ec2032..cc4cdee 100644
--- a/training/util.py
+++ b/training/util.py
@@ -41,14 +41,16 @@ class AverageMeter:
 class CheckpointerBase:
     def __init__(
         self,
-        datamodule,
+        train_dataloader,
+        val_dataloader,
         output_dir: Path,
         sample_image_size: int,
         sample_batches: int,
         sample_batch_size: int,
         seed: Optional[int] = None
     ):
-        self.datamodule = datamodule
+        self.train_dataloader = train_dataloader
+        self.val_dataloader = val_dataloader
         self.output_dir = output_dir
         self.sample_image_size = sample_image_size
         self.seed = seed if seed is not None else torch.random.seed()
@@ -70,15 +72,16 @@ class CheckpointerBase:
     ):
         samples_path = Path(self.output_dir).joinpath("samples")
 
-        train_data = self.datamodule.train_dataloader
-        val_data = self.datamodule.val_dataloader
-
         generator = torch.Generator(device=pipeline.device).manual_seed(self.seed)
 
         grid_cols = min(self.sample_batch_size, 4)
         grid_rows = (self.sample_batches * self.sample_batch_size) // grid_cols
 
-        for pool, data, gen in [("stable", val_data, generator), ("val", val_data, None), ("train", train_data, None)]:
+        for pool, data, gen in [
+            ("stable", self.val_dataloader, generator),
+            ("val", self.val_dataloader, None),
+            ("train", self.train_dataloader, None)
+        ]:
             all_samples = []
             file_path = samples_path.joinpath(pool, f"step_{step}.jpg")
             file_path.parent.mkdir(parents=True, exist_ok=True)