Fix OOM
Browse files- models/vallex.py +1 -0
models/vallex.py
CHANGED
|
@@ -676,6 +676,7 @@ class VALLE(VALLF):
|
|
| 676 |
y_emb[:, prefix_len:] += embedding_layer(samples)
|
| 677 |
|
| 678 |
assert len(codes) == self.num_quantizers
|
|
|
|
| 679 |
gc.collect()
|
| 680 |
return torch.stack(codes, dim=-1)
|
| 681 |
|
|
|
|
| 676 |
y_emb[:, prefix_len:] += embedding_layer(samples)
|
| 677 |
|
| 678 |
assert len(codes) == self.num_quantizers
|
| 679 |
+
del text_language_id, prompt_language_id, y_emb, x, y_pos, xy_pos, xy_dec, logits, samples
|
| 680 |
gc.collect()
|
| 681 |
return torch.stack(codes, dim=-1)
|
| 682 |
|