Update min_dalle/min_dalle.py

min_dalle/min_dalle.py

@@ -239,9 +239,9 @@ 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.cpu.amp.autocast(dtype=torch.Bfloat16):
+            with torch.cuda.amp.autocast(dtype=self.bfloat16):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
                     settings=settings,
                     attention_mask=attention_mask,
@@ -250,9 +250,8 @@ class MinDalle:
                     prev_tokens=image_tokens[i],
                     token_index=token_indices[[i]]
                 )
-       
 
-            with torch.cpu.amp.autocast(dtype=torch.Bfloat16):
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                 if ((i + 1) % 32 == 0 and progressive_outputs) or i + 1 == 256:
                     yield self.image_grid_from_tokens(
                         image_tokens=image_tokens[1:].T,