Update imagic.py
Browse files
imagic.py
CHANGED
@@ -274,8 +274,8 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
274 |
|
275 |
with torch.no_grad():
|
276 |
torch.cuda.empty_cache()
|
277 |
-
|
278 |
-
|
279 |
|
280 |
# Checks if the accelerator has performed an optimization step behind the scenes
|
281 |
if accelerator.sync_gradients:
|
@@ -318,8 +318,8 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
318 |
|
319 |
with torch.no_grad():
|
320 |
torch.cuda.empty_cache()
|
321 |
-
|
322 |
-
|
323 |
|
324 |
# Checks if the accelerator has performed an optimization step behind the scenes
|
325 |
if accelerator.sync_gradients:
|
|
|
274 |
|
275 |
with torch.no_grad():
|
276 |
torch.cuda.empty_cache()
|
277 |
+
optimizer.step()
|
278 |
+
optimizer.zero_grad()
|
279 |
|
280 |
# Checks if the accelerator has performed an optimization step behind the scenes
|
281 |
if accelerator.sync_gradients:
|
|
|
318 |
|
319 |
with torch.no_grad():
|
320 |
torch.cuda.empty_cache()
|
321 |
+
optimizer.step()
|
322 |
+
optimizer.zero_grad()
|
323 |
|
324 |
# Checks if the accelerator has performed an optimization step behind the scenes
|
325 |
if accelerator.sync_gradients:
|