diff options
Diffstat (limited to 'train_dreambooth.py')
-rw-r--r-- | train_dreambooth.py | 2 |
1 files changed, 1 insertions, 1 deletions
diff --git a/train_dreambooth.py b/train_dreambooth.py index 2f913e7..1a79b2b 100644 --- a/train_dreambooth.py +++ b/train_dreambooth.py | |||
@@ -687,7 +687,7 @@ def main(): | |||
687 | ).to(accelerator.device) | 687 | ).to(accelerator.device) |
688 | pipeline.set_progress_bar_config(dynamic_ncols=True) | 688 | pipeline.set_progress_bar_config(dynamic_ncols=True) |
689 | 689 | ||
690 | with torch.autocast("cuda"), torch.inference_mode(): | 690 | with torch.inference_mode(): |
691 | for batch in batched_data: | 691 | for batch in batched_data: |
692 | image_name = [item.class_image_path for item in batch] | 692 | image_name = [item.class_image_path for item in batch] |
693 | prompt = [item.cprompt for item in batch] | 693 | prompt = [item.cprompt for item in batch] |