Update imagic.py

imagic.py

@@ -272,8 +272,10 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
                 loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
                 accelerator.backward(loss)
 
-                optimizer.step()
-                optimizer.zero_grad()
+                with torch.no_grad()
+                    torch.cuda.empty_cache()
+                    optimizer.step()
+                    optimizer.zero_grad()
 
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients:
@@ -314,7 +316,8 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
                 loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
                 accelerator.backward(loss)
 
-                with torch.no_grad():
+                with torch.no_grad()
+                    torch.cuda.empty_cache()
                     optimizer.step()
                     optimizer.zero_grad()