Fix memory leak and reduce memory usage

This commit is contained in:
Jairo Correa
2022-09-28 22:14:13 -03:00
parent 041d2aefc0
commit c938679de7
6 changed files with 42 additions and 16 deletions

View File

@@ -1,4 +1,5 @@
import torch
import gc
# has_mps is only available in nightly pytorch (for now), `getattr` for compatibility
from modules import errors
@@ -17,8 +18,8 @@ def get_optimal_device():
return cpu
def torch_gc():
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
torch.cuda.ipc_collect()