Update imagic.py

imagic.py

@@ -314,8 +314,9 @@ 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():
+                    optimizer.step()
+                    optimizer.zero_grad()
 
             # Checks if the accelerator has performed an optimization step behind the scenes
             if accelerator.sync_gradients: