Spaces:

fffiloni
/

imagic-stable-diffusion

Paused

fffiloni commited on Nov 30, 2022

Commit

b7f1d3d

•

1 Parent(s): e720880

Update imagic.py

Files changed (1) hide show

imagic.py CHANGED Viewed

@@ -272,7 +272,7 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
                 loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
                 accelerator.backward(loss)
-                with torch.no_grad()
                     torch.cuda.empty_cache()
                     optimizer.step()
                     optimizer.zero_grad()
@@ -316,7 +316,7 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
                 loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
                 accelerator.backward(loss)
-                with torch.no_grad()
                     torch.cuda.empty_cache()
                     optimizer.step()
                     optimizer.zero_grad()

                 loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
                 accelerator.backward(loss)
+                with torch.no_grad():
                     torch.cuda.empty_cache()
                     optimizer.step()
                     optimizer.zero_grad()
                 loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
                 accelerator.backward(loss)
+                with torch.no_grad():
                     torch.cuda.empty_cache()
                     optimizer.step()
                     optimizer.zero_grad()