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

34
torch_lazy_loader.py
View File

@ -49,7 +49,7 @@ class _LazyUnpickler(pickle.Unpickler):
if key not in self.lazy_loaded_storages: if key not in self.lazy_loaded_storages:
self.lazy_loaded_storages[key] = LazyTensor(storage_type, key, location, nelements) self.lazy_loaded_storages[key] = LazyTensor(storage_type, key, location, nelements)
return self.lazy_loaded_storages[key] return self.lazy_loaded_storages[key]
def load(self, *args, **kwargs): def load(self, *args, **kwargs):
@ -72,24 +72,26 @@ def use_lazy_torch_load(enable=True, callback: Optional[Callable] = None):
yield False yield False
return return
old_unpickler = pickle.Unpickler try:
pickle.Unpickler = _LazyUnpickler old_unpickler = pickle.Unpickler
pickle.Unpickler = _LazyUnpickler
old_rebuild_tensor = torch._utils._rebuild_tensor old_rebuild_tensor = torch._utils._rebuild_tensor
torch._utils._rebuild_tensor = _rebuild_tensor torch._utils._rebuild_tensor = _rebuild_tensor
old_torch_load = torch.load old_torch_load = torch.load
def torch_load(f, map_location=None, pickle_module=pickle, **pickle_load_args): def torch_load(f, map_location=None, pickle_module=pickle, **pickle_load_args):
retval = old_torch_load(f=f, map_location=map_location, pickle_module=pickle_module, **pickle_load_args) retval = old_torch_load(f=f, map_location=map_location, pickle_module=pickle_module, **pickle_load_args)
if callback is not None: if callback is not None:
callback(retval, f=f, map_location=map_location, pickle_module=pickle_module, **pickle_load_args) callback(retval, f=f, map_location=map_location, pickle_module=pickle_module, **pickle_load_args)
return retval return retval
torch.load = torch_load torch.load = torch_load
yield True yield True
pickle.Unpickler = old_unpickler finally:
torch._utils._rebuild_tensor = old_rebuild_tensor pickle.Unpickler = old_unpickler
torch.load = old_torch_load torch._utils._rebuild_tensor = old_rebuild_tensor
torch.load = old_torch_load