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

Add PyTorch 1.11 support for lazy loader

Browse Source
This commit is contained in:
Gnome Ann
2022-03-17 12:51:41 -04:00
parent 9235754eb9
commit eaf190469d
4 changed files with 27 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tpu_mtj_backend.py
View File

@@ -961,7 +961,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
# the least possible memory usage, we create them as meta
# tensors, which don't take up any actual CPU or TPU memory.
if key not in model_spec:
model_dict[key] = torch.empty(model_dict[key].shape, dtype=model_dict[key].storage_type(0).dtype, device="meta")
model_dict[key] = torch.empty(model_dict[key].shape, dtype=model_dict[key].dtype, device="meta")
continue
storage_key = model_dict[key].key