Skip to content

GitLab

"tools/git@developer.sourcefind.cn:OpenDAS/nni.git" did not exist on "cd3a912a07c2187fcd9e81205da35b97746b76f3"
Commit d63705d9 authored by 藍+85CD's avatar 藍+85CD
Browse files

Support releases all unoccupied cached memory from XPU

parent 04d9bc13
Hide whitespace changes
Inline Side-by-side
Showing with 4 additions and 0 deletions
execution.py
View file @ d63705d9
...@@ -10,6 +10,8 @@ import gc ...@@ -10,6 +10,8 @@ import gc
import torch import torch
import nodes import nodes
from model_management import xpu_available
def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}): def get_input_data(inputs, class_def, unique_id, outputs={}, prompt={}, extra_data={}):
valid_inputs = class_def.INPUT_TYPES() valid_inputs = class_def.INPUT_TYPES()
input_data_all = {} input_data_all = {}
...@@ -206,6 +208,8 @@ class PromptExecutor: ...@@ -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 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.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
elif xpu_available:
torch.xpu.empty_cache()
def validate_inputs(prompt, item): def validate_inputs(prompt, item):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment