Update min_dalle/models/dalle_bart_decoder.py

min_dalle/models/dalle_bart_decoder.py

@@ -3,6 +3,7 @@ import torch
 from torch import nn, LongTensor, FloatTensor, BoolTensor
 from .dalle_bart_encoder import GLU, AttentionBase
 import gc
+import tracemalloc
 
 IMAGE_TOKEN_COUNT = 256
 
@@ -154,6 +155,12 @@ class DalleBartDecoder(nn.Module):
         decoder_state += self.embed_positions.forward(token_index_batched)
         decoder_state = self.layernorm_embedding.forward(decoder_state)
         decoder_state = decoder_state[:, None]
+        
+        tracemalloc.start()
+        print("--")
+        # displaying the memory
+        print(tracemalloc.get_traced_memory())
+        
         for i in range(self.layer_count):
             decoder_state, attention_state[i] = self.layers[i].forward(
                 decoder_state,
@@ -173,6 +180,7 @@ class DalleBartDecoder(nn.Module):
             logits[:image_count] * (1 - supercondition_factor) + 
             logits[image_count:] * supercondition_factor
         )
+        print(tracemalloc.get_traced_memory())
         del supercondition_factor 
         logits_sorted, _ = logits.sort(descending=True)
         is_kept = logits >= logits_sorted[:, top_k - 1]
@@ -188,4 +196,6 @@ class DalleBartDecoder(nn.Module):
         del logits
         gc.collect()
         
+        print(tracemalloc.get_traced_memory())
+        
         return image_tokens, attention_state