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

Slight fix

Browse Source
This commit is contained in:
ebolam
2022-10-24 19:01:11 -04:00
parent 47c832fde4
commit 46ac4ec009
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tpu_mtj_backend.py
View File

@@ -1015,7 +1015,7 @@ def read_neox_checkpoint(state, path, config, checkpoint_shards=2):
if socketio is None:
bar = tqdm(total=tqdm_length, desc="Loading from NeoX checkpoint")
else:
bar = tqdm(total=tqdm_length, desc="Loading from NeoX checkpoint", file=Send_to_socketio)
bar = tqdm(total=tqdm_length, desc="Loading from NeoX checkpoint", file=Send_to_socketio())
koboldai_vars.status_message = "Loading TPU"
koboldai_vars.total_layers = tqdm_length
koboldai_vars.loaded_layers = 0
@@ -1296,7 +1296,7 @@ def load_model(path: str, driver_version="tpu_driver0.1_dev20210607", hf_checkpo
if socketio is None:
utils.bar = tqdm(total=num_tensors, desc="Loading model tensors")
else:
utils.bar = tqdm(total=num_tensors, desc="Loading model tensors", file=Send_to_socketio)
utils.bar = tqdm(total=num_tensors, desc="Loading model tensors", file=Send_to_socketio())
koboldai_vars.status_message = "Loading model"
koboldai_vars.loaded_layers = 0
koboldai_vars.total_layers = num_tensors