Spaces:

JMalott
/

ai_architecture

Runtime error

JMalott commited on Sep 6, 2022

Commit

cf33d2b

1 Parent(s): d1d5dcf

Update min_dalle/min_dalle.py

Files changed (1) hide show

min_dalle/min_dalle.py CHANGED Viewed

@@ -241,7 +241,7 @@ class MinDalle:
             torch.cuda.empty_cache()
             #torch.cpu.empty_cache()
-            with torch.cuda.amp.autocast(dtype=self.bfloat16):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
                     settings=settings,
                     attention_mask=attention_mask,

             torch.cuda.empty_cache()
             #torch.cpu.empty_cache()
+            with torch.cuda.amp.autocast(dtype=torch.bfloat16):
                 image_tokens[i + 1], attention_state = self.decoder.forward(
                     settings=settings,
                     attention_mask=attention_mask,