From 6473aa1dd7cd927852957b977cfdf4f86b4f17b9 Mon Sep 17 00:00:00 2001
From: Isotr0py <2037008807@qq.com>
Date: Wed, 8 Feb 2023 21:32:21 +0800
Subject: [PATCH] fix Input type error in noise_pred when using DDP

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

diff --git a/train_network.py b/train_network.py
index fc387bc3..59f74211 100644
--- a/train_network.py
+++ b/train_network.py
@@ -417,7 +417,8 @@ def train(args):
 
         # Add noise to the latents according to the noise magnitude at each timestep
         # (this is the forward diffusion process)
-        noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
+        with autocast():
+          noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # Predict the noise residual
         noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample