From 509e4e3f9627f674a2df6153639d8f5d2b41f55b Mon Sep 17 00:00:00 2001 From: patil-suraj Date: Mon, 26 Sep 2022 14:06:18 +0200 Subject: [PATCH] actuallu fix saving --- examples/dreambooth/train_dreambooth.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/dreambooth/train_dreambooth.py b/examples/dreambooth/train_dreambooth.py index 9767adcf5d..e615fe5416 100644 --- a/examples/dreambooth/train_dreambooth.py +++ b/examples/dreambooth/train_dreambooth.py @@ -366,8 +366,8 @@ def main(): with context: images = pipeline(example["prompt"]).images - for image in images: - image.save(class_images_dir / f"{example['index'] + cur_class_images}.jpg") + for image, index in (images, example["index"]): + image.save(class_images_dir / f"{index + cur_class_images}.jpg") del pipeline if torch.cuda.is_available():