Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -242,7 +242,7 @@ class MinDalle:
 
             torch.cuda.empty_cache()                
             
-            time.sleep(0.5)
+            #time.sleep(0.5)
             
             with torch.cuda.amp.autocast(dtype=self.dtype):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
@@ -253,6 +253,8 @@ class MinDalle:
                     prev_tokens=image_tokens[i],
                     token_index=token_indices[[i]]
                 )
+                
+            torch.cuda.empty_cache()
 
             with torch.cuda.amp.autocast(dtype=torch.float16):
                 if ((i + 1) % 16 == 0 and progressive_outputs) or i + 1 == 256: