This commit is contained in:
wfjsw 2024-02-20 21:08:07 -06:00
parent fffc902698
commit 8828c9ecc5

View File

@ -1,6 +1,5 @@
from contextlib import contextmanager, nullcontext from contextlib import contextmanager, nullcontext
import torch import torch
from torch.utils.weak import WeakIdKeyDictionary
from modules import devices, shared, patches from modules import devices, shared, patches
module_in_gpu = None module_in_gpu = None
@ -86,7 +85,7 @@ class RTTensorMoverPatches:
self.stash[id(after_calc_event)] = (weight, bias, after_calc_event) self.stash[id(after_calc_event)] = (weight, bias, after_calc_event)
to_remove = [] to_remove = []
for k, (w, b, e) in self.stash.items(): for k, (_, _, e) in self.stash.items():
if e.query(): if e.query():
to_remove.append(k) to_remove.append(k)