Update imagic.py

imagic.py

@@ -272,8 +272,6 @@ 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()
 
@@ -291,8 +289,6 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
         text_embeddings.requires_grad_(False)
 
         # Now we fine tune the unet to better reconstruct the image
-        with torch.no_grad():
-            torch.cuda.empty_cache()
         self.unet.requires_grad_(True)
         self.unet.train()
         optimizer = torch.optim.Adam(
@@ -317,12 +313,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():
-                    torch.cuda.empty_cache()
-
-                with torch.no_grad():
-                    optimizer.step()
+               
+                optimizer.step()
                 optimizer.zero_grad()
 
             # Checks if the accelerator has performed an optimization step behind the scenes