Update min_dalle/models/dalle_bart_decoder.py

min_dalle/models/dalle_bart_decoder.py

@@ -175,7 +175,7 @@ class DalleBartDecoder(nn.Module):
         logits /= temperature
         logits.exp_()
         logits *= is_kept.to(torch.float32)
-        image_tokens = torch.multinomial(logits, 1)[:, 0]
-        del logits
-        print("hi")
-        return image_tokens, attention_state
+        #image_tokens = torch.multinomial(logits, 1)[:, 0]
+        #del logits
+        
+        return torch.multinomial(logits, 1)[:, 0], attention_state