Update min_dalle/models/dalle_bart_decoder.py

min_dalle/models/dalle_bart_decoder.py

@@ -161,6 +161,7 @@ class DalleBartDecoder(nn.Module):
             )
         decoder_state = self.final_ln(decoder_state)
         logits = self.lm_head(decoder_state)
+        del decorder_state
         temperature = settings[[0]]
         top_k = settings[[1]].to(torch.long)
         supercondition_factor = settings[[2]]
@@ -171,11 +172,13 @@ class DalleBartDecoder(nn.Module):
         )
         logits_sorted, _ = logits.sort(descending=True)
         is_kept = logits >= logits_sorted[:, top_k - 1]
+        del top_k
         logits -= logits_sorted[:, [0]]
+        del logits_sorted
         logits /= temperature
         logits.exp_()
         logits *= is_kept.to(torch.float32)
-        #image_tokens = torch.multinomial(logits, 1)[:, 0]
-        #del logits
+        image_tokens = torch.multinomial(logits, 1)[:, 0]
+        del logits
         
-        return torch.multinomial(logits, 1)[:, 0], attention_state
+        return image_tokens, attention_state