Support releases all unoccupied cached memory from XPU

This commit is contained in:
藍+85CD 2023-04-15 15:50:51 +08:00
parent 04d9bc13af
commit d63705d919
1 changed files with 4 additions and 0 deletions

View File

@ -10,6 +10,8 @@ import gc
import torch
import nodes
from model_management import xpu_available
def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):
valid_inputs = class_def.INPUT_TYPES()
input_data_all = {}
@ -206,6 +208,8 @@ class PromptExecutor:
if torch.version.cuda: #This seems to make things worse on ROCm so I only do it for cuda
torch.cuda.empty_cache()
torch.cuda.ipc_collect()
elif xpu_available:
torch.xpu.empty_cache()
def validate_inputs(prompt, item):