From 193fb1660a4be43edeb2d121b5fa49a33a437f86 Mon Sep 17 00:00:00 2001 From: oobabooga <112222186+oobabooga@users.noreply.github.com> Date: Wed, 22 Feb 2023 00:35:10 -0300 Subject: [PATCH] Conversion seems to work better this way --- convert-to-flexgen.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/convert-to-flexgen.py b/convert-to-flexgen.py index e59eb99..18afa9b 100644 --- a/convert-to-flexgen.py +++ b/convert-to-flexgen.py @@ -45,9 +45,9 @@ if __name__ == '__main__': model_name = path.name print(f"Loading {model_name}...") - disable_torch_init() - model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16, _fast_init=True) - restore_torch_init() + #disable_torch_init() + model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16, low_cpu_mem_usage=True) + #restore_torch_init() tokenizer = AutoTokenizer.from_pretrained(path)