TPU Fixes

This commit is contained in:
ebolam
2023-05-24 18:55:31 -04:00
parent f0f646ae7b
commit 5fe8c71b2e
2 changed files with 5 additions and 4 deletions

View File

@@ -186,6 +186,7 @@ class model_backend(HFInferenceModel):
tpu_mtj_backend.load_model(
utils.koboldai_vars.model,
self.model_type,
hf_checkpoint=utils.koboldai_vars.model
not in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX")
and utils.koboldai_vars.use_colab_tpu,
@@ -202,7 +203,7 @@ class model_backend(HFInferenceModel):
if (
utils.koboldai_vars.badwordsids is koboldai_settings.badwordsids_default
and utils.koboldai_vars.model_type not in ("gpt2", "gpt_neo", "gptj")
and self.model_type not in ("gpt2", "gpt_neo", "gptj")
):
utils.koboldai_vars.badwordsids = [
[v]