Seich
/
mayavoz
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

empty cache

This commit is contained in:
shahules786 2022-10-10 21:03:39 +05:30
parent 4363821c75
commit c193a48e8e
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
enhancer/models/model.py
View File

@ -110,6 +110,7 @@ class Model(pl.LightningModule):
def setup(self, stage: Optional[str] = None): def setup(self, stage: Optional[str] = None):
if stage == "fit": if stage == "fit":
torch.cuda.empty_cache()
self.dataset.setup(stage) self.dataset.setup(stage)
self.dataset.model = self self.dataset.model = self