From 46ac4ec0097af072e5c3e66ba0267e60e1163ffa Mon Sep 17 00:00:00 2001 From: ebolam Date: Mon, 24 Oct 2022 19:01:11 -0400 Subject: [PATCH] Slight fix --- tpu_mtj_backend.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tpu_mtj_backend.py b/tpu_mtj_backend.py index aaa22154..492102a8 100644 --- a/tpu_mtj_backend.py +++ b/tpu_mtj_backend.py @@ -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