let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

optimize mem

Browse Source
This commit is contained in:
Qing
2024-01-09 22:54:20 +08:00
parent e94a94e3c2
commit db4a6f4547
3 changed files with 18 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
iopaint/__init__.py
View File

@@ -1,6 +1,14 @@
import os
os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"
# https://github.com/pytorch/pytorch/issues/27971#issuecomment-1768868068
os.environ["ONEDNN_PRIMITIVE_CACHE_CAPACITY"] = "1"
os.environ["LRU_CACHE_CAPACITY"] = "1"
# prevent CPU memory leak when run model on GPU
# https://github.com/pytorch/pytorch/issues/98688#issuecomment-1869288431
# https://github.com/pytorch/pytorch/issues/108334#issuecomment-1752763633
os.environ["TORCH_CUDNN_V8_API_LRU_CACHE_LIMIT"] = "1"
import warnings