Update min_dalle/models/dalle_bart_decoder.py

min_dalle/models/dalle_bart_decoder.py

@@ -162,14 +162,13 @@ class DalleBartDecoder(nn.Module):
         print(tracemalloc.get_traced_memory())
         
         for i in range(self.layer_count):
-            decoder_state, attention_state[i] = self.layers[i].forward(
+            del decoder_state, attention_state[i] = self.layers[i].forward(
                 decoder_state,
                 encoder_state,
                 attention_state[i],
                 attention_mask,
                 token_index
             )
-            del decoder_state
         print(tracemalloc.get_traced_memory())
         decoder_state = self.final_ln(decoder_state)
         logits = self.lm_head(decoder_state)