Conversion seems to work better this way

2023-02-22 00:35:10 -03:00
parent eef6fc3cbf
commit 193fb1660a
1 changed files with 3 additions and 3 deletions
--- 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)