Better dataset prompt handling

1 files changed, 1 insertions, 1 deletions

diff --git a/train_ti.py b/train_ti.py
index e272b5d..cc208f0 100644
--- a/train_ti.py
+++ b/train_ti.py
@@ -621,7 +621,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]