1 files changed, 2 insertions, 2 deletions

diff --git a/training/util.py b/training/util.py
index bc466e2..6f42228 100644
--- a/training/util.py
+++ b/training/util.py
@@ -58,8 +58,8 @@ class CheckpointerBase:
     def save_samples(self, pipeline, step, num_inference_steps, guidance_scale=7.5, eta=0.0):
         samples_path = Path(self.output_dir).joinpath("samples")
 
-        train_data = self.datamodule.train_dataloader()
-        val_data = self.datamodule.val_dataloader()
+        train_data = self.datamodule.train_dataloaders[0]
+        val_data = self.datamodule.val_dataloader
 
         generator = torch.Generator(device=pipeline.device).manual_seed(self.seed)