Update ProTalk_ModelExporter.py
Browse files- ProTalk_ModelExporter.py +2 -2
ProTalk_ModelExporter.py
CHANGED
|
@@ -1,6 +1,6 @@
|
|
| 1 |
from transformers import AutoModelForCausalLM
|
| 2 |
from safetensors.torch import save_file
|
| 3 |
|
| 4 |
-
|
| 5 |
-
model = AutoModelForCausalLM.from_pretrained(
|
| 6 |
save_file(model.state_dict(), "ProTalkModel.safetensors")
|
|
|
|
| 1 |
from transformers import AutoModelForCausalLM
|
| 2 |
from safetensors.torch import save_file
|
| 3 |
|
| 4 |
+
model_name = "microsoft/phi-2"
|
| 5 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
|
| 6 |
save_file(model.state_dict(), "ProTalkModel.safetensors")
|