From 9a9ac79edff44e9c0cdb28d127f62fe5ce0cca07 Mon Sep 17 00:00:00 2001 From: Isotr0py <2037008807@qq.com> Date: Wed, 8 Feb 2023 22:30:20 +0800 Subject: [PATCH] correct wrong inserted code for noise_pred fix --- train_network.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/train_network.py b/train_network.py index 59f74211..f247c74e 100644 --- a/train_network.py +++ b/train_network.py @@ -417,11 +417,11 @@ 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) + noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps) # Predict the noise residual - noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample + with autocast(): + noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample if args.v_parameterization: # v-parameterization training