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:
|