Update gligen/ldm/models/diffusion/plms.py

gligen/ldm/models/diffusion/plms.py

@@ -189,6 +189,7 @@ class PLMSSampler(object):
         
         print("optimize", index1)
         self.model.train()
+        torch.cuda.empty_cache()
         while loss.item() > loss_threshold and iteration < max_iter and (index1 < max_index) :
             print('iter', iteration)
             # import pdb; pdb.set_trace()
@@ -218,7 +219,7 @@ class PLMSSampler(object):
             x = x.detach()
             iteration += 1
             torch.cuda.empty_cache()
-            
+            torch.cuda.synchronize()
         
         return x