From a4d2f7545b071445c0cfada70df7b12f119afa16 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Wed, 4 Sep 2024 19:42:30 +0300
Subject: [PATCH] Update flux_train_utils.py

---
 library/flux_train_utils.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/library/flux_train_utils.py b/library/flux_train_utils.py
index f67ed6d..de65d38 100644
--- a/library/flux_train_utils.py
+++ b/library/flux_train_utils.py
@@ -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(