Update min_dalle/models/dalle_bart_decoder.py

min_dalle/models/dalle_bart_decoder.py

@@ -2,6 +2,7 @@ from typing import Tuple, List
 import torch
 from torch import nn, LongTensor, FloatTensor, BoolTensor
 from .dalle_bart_encoder import GLU, AttentionBase
+import gc
 
 IMAGE_TOKEN_COUNT = 256
 
@@ -100,6 +101,8 @@ class DecoderLayer(nn.Module):
         decoder_state = self.glu.forward(decoder_state)
         decoder_state = residual + decoder_state
 
+        
+
         return decoder_state, attention_state
 
 
@@ -170,6 +173,7 @@ class DalleBartDecoder(nn.Module):
             logits[:image_count] * (1 - supercondition_factor) + 
             logits[image_count:] * supercondition_factor
         )
+        del supercondition_factor 
         logits_sorted, _ = logits.sort(descending=True)
         is_kept = logits >= logits_sorted[:, top_k - 1]
         del top_k
@@ -179,7 +183,9 @@ class DalleBartDecoder(nn.Module):
         del temperature
         logits.exp_()
         logits *= is_kept.to(torch.float32)
+        del is_kept
         image_tokens = torch.multinomial(logits, 1)[:, 0]
         del logits
+        gc.collect()
         
         return image_tokens, attention_state