Update min_dalle/models/dalle_bart_decoder.py

min_dalle/models/dalle_bart_decoder.py

@@ -169,11 +169,14 @@ class DalleBartDecoder(nn.Module):
                 attention_mask,
                 token_index
             )
+        print(tracemalloc.get_traced_memory())
         decoder_state = self.final_ln(decoder_state)
         logits = self.lm_head(decoder_state)
+        print(tracemalloc.get_traced_memory())
         del decoder_state 
         temperature = settings[[0]]
         top_k = settings[[1]].to(torch.long)
+        print(tracemalloc.get_traced_memory())
         supercondition_factor = settings[[2]]
         logits = logits[:, -1, : 2 ** 14]
         logits: FloatTensor = (