diff options
Diffstat (limited to 'training')
-rw-r--r-- | training/util.py | 2 |
1 files changed, 1 insertions, 1 deletions
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: | |||
59 | def save_samples(self, pipeline, step, num_inference_steps, guidance_scale=7.5, eta=0.0): | 59 | def save_samples(self, pipeline, step, num_inference_steps, guidance_scale=7.5, eta=0.0): |
60 | samples_path = Path(self.output_dir).joinpath("samples") | 60 | samples_path = Path(self.output_dir).joinpath("samples") |
61 | 61 | ||
62 | train_data = self.datamodule.train_dataloaders[0] | 62 | train_data = self.datamodule.train_dataloader |
63 | val_data = self.datamodule.val_dataloader | 63 | val_data = self.datamodule.val_dataloader |
64 | 64 | ||
65 | generator = torch.Generator(device=pipeline.device).manual_seed(self.seed) | 65 | generator = torch.Generator(device=pipeline.device).manual_seed(self.seed) |