Update audio_litmodule.py

JusperLee · Sep 23, 2024 · 4630082 · 4630082
1 parent e036588
commit 4630082
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/look2hear/system/audio_litmodule.py b/look2hear/system/audio_litmodule.py
@@ -158,6 +158,7 @@ def on_validation_epoch_end(self):
         )
 
         self.validation_step_outputs.clear()  # free memory
+        torch.cuda.empty_cache()
 
     def test_step(self, batch, batch_nb):
         mixtures, targets = batch