mirror of
https://git.datalinker.icu/comfyanonymous/ComfyUI
synced 2026-08-01 22:18:51 +08:00
Merge e148a7f1bb6ecd9f756526da0a93ac45dc78234c into b2ef58e2b17e73ca8cd376a1cdc976518ebbc168
This commit is contained in:
commit
6eb43de4b1
@ -1098,6 +1098,15 @@ if not args.disable_pinned_memory:
|
||||
MAX_PINNED_MEMORY = get_total_memory(torch.device("cpu")) * 0.95
|
||||
logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
|
||||
|
||||
def discard_cuda_async_error():
|
||||
try:
|
||||
a = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
|
||||
b = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
|
||||
_ = a + b
|
||||
torch.cuda.synchronize()
|
||||
except torch.AcceleratorError:
|
||||
#Dump it! We already know about it from the synchronous return
|
||||
pass
|
||||
|
||||
def pin_memory(tensor):
|
||||
global TOTAL_PINNED_MEMORY
|
||||
@ -1128,6 +1137,8 @@ def pin_memory(tensor):
|
||||
PINNED_MEMORY[ptr] = size
|
||||
TOTAL_PINNED_MEMORY += size
|
||||
return True
|
||||
else:
|
||||
discard_cuda_async_error()
|
||||
|
||||
return False
|
||||
|
||||
@ -1156,6 +1167,8 @@ def unpin_memory(tensor):
|
||||
if len(PINNED_MEMORY) == 0:
|
||||
TOTAL_PINNED_MEMORY = 0
|
||||
return True
|
||||
else:
|
||||
discard_cuda_async_error()
|
||||
|
||||
return False
|
||||
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user