Przeglądaj źródła

Conversion seems to work better this way

oobabooga 2 lat temu
rodzic
commit
193fb1660a
1 zmienionych plików z 3 dodań i 3 usunięć
  1. 3 3
      convert-to-flexgen.py

+ 3 - 3
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)