Spaces:
Runtime error
Runtime error
Update min_dalle/models/dalle_bart_decoder.py
Browse files
min_dalle/models/dalle_bart_decoder.py
CHANGED
|
@@ -162,7 +162,7 @@ class DalleBartDecoder(nn.Module):
|
|
| 162 |
print(tracemalloc.get_traced_memory())
|
| 163 |
|
| 164 |
for i in range(self.layer_count):
|
| 165 |
-
|
| 166 |
decoder_state,
|
| 167 |
encoder_state,
|
| 168 |
attention_state[i],
|
|
|
|
| 162 |
print(tracemalloc.get_traced_memory())
|
| 163 |
|
| 164 |
for i in range(self.layer_count):
|
| 165 |
+
decoder_state, attention_state[i] = self.layers[i].forward(
|
| 166 |
decoder_state,
|
| 167 |
encoder_state,
|
| 168 |
attention_state[i],
|