tanner/min-dalle-test
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Merge pull request #57 from w4ffl35/bugfix-for-mega-model

Browse Source 
Fixes #16 - mega model running out of memory
This commit is contained in:
Brett Kuprel 2022-07-04 07:16:39 -04:00 committed by GitHub
commit 9602055eaa
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
min_dalle/min_dalle.py
View File

@ -142,6 +142,7 @@ class MinDalle:
params = torch.load(self.detoker_params_path) params = torch.load(self.detoker_params_path)
self.detokenizer.load_state_dict(params) self.detokenizer.load_state_dict(params)
del params del params
torch.cuda.empty_cache()
if torch.cuda.is_available(): self.detokenizer = self.detokenizer.cuda() if torch.cuda.is_available(): self.detokenizer = self.detokenizer.cuda()
@ -175,6 +176,7 @@ class MinDalle:
encoder_state encoder_state
) )
if not self.is_reusable: del self.decoder if not self.is_reusable: del self.decoder
torch.cuda.empty_cache()
return image_tokens return image_tokens