Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -239,17 +239,17 @@ 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,
-                    attention_mask=attention_mask,
-                    encoder_state=encoder_state,
-                    attention_state=attention_state,
-                    prev_tokens=image_tokens[i],
-                    token_index=token_indices[[i]]
-                )
+            #with torch.cuda.amp.autocast(dtype=self.dtype):
+            image_tokens[i + 1], attention_state = self.decoder.forward(
+                settings=settings,
+                attention_mask=attention_mask,
+                encoder_state=encoder_state,
+                attention_state=attention_state,
+                prev_tokens=image_tokens[i],
+                token_index=token_indices[[i]]
+            )
 
            # with torch.cuda.amp.autocast(dtype=torch.float32):
             if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 256: