empty cache

2022-10-10 21:03:39 +05:30 · 2022-10-10 21:03:39 +05:30 · c193a48e8e
parent 4363821c75
commit c193a48e8e
1 changed files with 1 additions and 0 deletions
--- a/enhancer/models/model.py
+++ b/enhancer/models/model.py
@ -110,6 +110,7 @@ class Model(pl.LightningModule):

    def setup(self, stage: Optional[str] = None):
        if stage == "fit":
+            torch.cuda.empty_cache()
            self.dataset.setup(stage)
            self.dataset.model = self