Dreambooth: Generate specialized class images from input prompts

1 files changed, 2 insertions, 4 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index 5fc2338..4c4da29 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -694,8 +694,7 @@ def main():
                     # Predict the noise residual
                     noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
 
-                    with accelerator.autocast():
-                        loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
+                    loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
 
                     accelerator.backward(loss)
 
@@ -766,8 +765,7 @@ def main():
 
                     noise_pred, noise = accelerator.gather_for_metrics((noise_pred, noise))
 
-                    with accelerator.autocast():
-                        loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
+                    loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
 
                     loss = loss.detach().item()
                     val_loss += loss