From f403ac61327444a37b8e503dd4bdb7a534f80783 Mon Sep 17 00:00:00 2001
From: Kohya S <ykumeykume@gmail.com>
Date: Thu, 23 Feb 2023 20:56:41 +0900
Subject: [PATCH] fix float32 training doesn't work in some case

---
 train_network.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/train_network.py b/train_network.py
index b01ec117..ee073ba8 100644
--- a/train_network.py
+++ b/train_network.py
@@ -361,7 +361,7 @@ def train(args):
         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # Predict the noise residual
-        with autocast():
+        with accelerator.autocast():
           noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
         if args.v_parameterization: