From 5571c4ebcb39813e2bd8585de30c64bb02f9d7fa Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Sun, 8 Jan 2023 09:43:22 +0100
Subject: Improved aspect ratio bucketing

---
 training/util.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

(limited to 'training/util.py')

diff --git a/training/util.py b/training/util.py
index 2b7f71d..ae6bfc4 100644
--- a/training/util.py
+++ b/training/util.py
@@ -59,7 +59,7 @@ 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_dataloaders[0]
+        train_data = self.datamodule.train_dataloader
         val_data = self.datamodule.val_dataloader
 
         generator = torch.Generator(device=pipeline.device).manual_seed(self.seed)
-- 
cgit v1.2.3-54-g00ecf