fffiloni commited on
Commit
f29e039
1 Parent(s): 44b9114

Update imagic.py

Browse files
Files changed (1) hide show
  1. imagic.py +2 -0
imagic.py CHANGED
@@ -291,6 +291,8 @@ class ImagicStableDiffusionPipeline(DiffusionPipeline):
291
  text_embeddings.requires_grad_(False)
292
 
293
  # Now we fine tune the unet to better reconstruct the image
 
 
294
  self.unet.requires_grad_(True)
295
  self.unet.train()
296
  optimizer = torch.optim.Adam(
 
291
  text_embeddings.requires_grad_(False)
292
 
293
  # Now we fine tune the unet to better reconstruct the image
294
+ with torch.no_grad():
295
+ torch.cuda.empty_cache()
296
  self.unet.requires_grad_(True)
297
  self.unet.train()
298
  optimizer = torch.optim.Adam(