Skip to content

Commit dfd616e

Browse files
authored
Avoid OOM when other tests are failing (#39758)
fix Co-authored-by: ydshieh <[email protected]>
1 parent 65df73a commit dfd616e

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

tests/models/musicgen_melody/test_modeling_musicgen_melody.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
)
3232
from transformers.testing_utils import (
3333
Expectations,
34+
cleanup,
3435
get_device_properties,
3536
is_torch_available,
3637
is_torchaudio_available,
@@ -1256,6 +1257,12 @@ class MusicgenMelodyIntegrationTests(unittest.TestCase):
12561257
def model(self):
12571258
return MusicgenMelodyForConditionalGeneration.from_pretrained("ylacombe/musicgen-melody").to(torch_device)
12581259

1260+
def setUp(self):
1261+
cleanup(torch_device, gc_collect=True)
1262+
1263+
def tearDown(self):
1264+
cleanup(torch_device, gc_collect=True)
1265+
12591266
@cached_property
12601267
def processor(self):
12611268
return MusicgenMelodyProcessor.from_pretrained("ylacombe/musicgen-melody")

0 commit comments

Comments
 (0)