rixty/KoboldAI-Client
1
1
Fork 0
mirror of https://github.com/KoboldAI/KoboldAI-Client.git synced 2025-06-05 21:59:24 +02:00
Code Issues Projects Releases Wiki Activity

Clean up when error is thrown in use_lazy_torch_load

Browse Source
This commit is contained in:
Gnome Ann
2022-03-01 19:30:22 -05:00
parent a0344b429c
commit 4fa4dbac50
1 changed files with 18 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
torch_lazy_loader.py
View File

@ -72,6 +72,7 @@ def use_lazy_torch_load(enable=True, callback: Optional[Callable] = None):
yield False yield False
return return
try:
old_unpickler = pickle.Unpickler old_unpickler = pickle.Unpickler
pickle.Unpickler = _LazyUnpickler pickle.Unpickler = _LazyUnpickler
@ -90,6 +91,7 @@ def use_lazy_torch_load(enable=True, callback: Optional[Callable] = None):
yield True yield True
finally:
pickle.Unpickler = old_unpickler pickle.Unpickler = old_unpickler
torch._utils._rebuild_tensor = old_rebuild_tensor torch._utils._rebuild_tensor = old_rebuild_tensor
torch.load = old_torch_load torch.load = old_torch_load