Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -239,7 +239,8 @@ class MinDalle:
                 break
             st.session_state.bar.progress(i/IMAGE_TOKEN_COUNT)
 
-            torch.cuda.empty_cache()                
+            torch.cuda.empty_cache()   
+            torch.cpu.empty_cache()              
             with torch.cuda.amp.autocast(dtype=self.dtype):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
                     settings=settings,