Update imagic.py
Browse files
imagic.py
CHANGED
@@ -272,7 +272,7 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
272 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
273 |
accelerator.backward(loss)
|
274 |
|
275 |
-
with torch.no_grad()
|
276 |
torch.cuda.empty_cache()
|
277 |
optimizer.step()
|
278 |
optimizer.zero_grad()
|
@@ -316,7 +316,7 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
|
|
316 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
317 |
accelerator.backward(loss)
|
318 |
|
319 |
-
with torch.no_grad()
|
320 |
torch.cuda.empty_cache()
|
321 |
optimizer.step()
|
322 |
optimizer.zero_grad()
|
|
|
272 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
273 |
accelerator.backward(loss)
|
274 |
|
275 |
+
with torch.no_grad():
|
276 |
torch.cuda.empty_cache()
|
277 |
optimizer.step()
|
278 |
optimizer.zero_grad()
|
|
|
316 |
loss = F.mse_loss(noise_pred, noise, reduction="none").mean([1, 2, 3]).mean()
|
317 |
accelerator.backward(loss)
|
318 |
|
319 |
+
with torch.no_grad():
|
320 |
torch.cuda.empty_cache()
|
321 |
optimizer.step()
|
322 |
optimizer.zero_grad()
|