Better dataset prompt handling

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():
             ).to(accelerator.device)
             pipeline.set_progress_bar_config(dynamic_ncols=True)
 
-            with torch.autocast("cuda"), torch.inference_mode():
+            with torch.inference_mode():
                 for batch in batched_data:
                     image_name = [item.class_image_path for item in batch]
                     prompt = [item.cprompt for item in batch]