Update flux_train_utils.py

2024-09-04 19:42:30 +03:00
parent 721dfc94f3
commit a4d2f7545b
1 changed files with 1 additions and 1 deletions
@@ -83,7 +83,7 @@ def sample_images(
    except Exception:
        pass

-    with torch.no_grad():
+    with torch.no_grad(), accelerator.autocast():
        image_tensor_list = []
        for prompt_dict in prompts:
            image_tensor = sample_image_inference(