fix Input type error in noise_pred when using DDP

train_network.py

@@ -417,6 +417,7 @@ def train(args):
 
         # Add noise to the latents according to the noise magnitude at each timestep
         # (this is the forward diffusion process)
+        with autocast():
           noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
         # Predict the noise residual