prelington commited on
Commit
4828ab6
·
verified ·
1 Parent(s): 773f5ca

Update ProTalk_ModelExporter.py

Browse files
Files changed (1) hide show
  1. 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
- model_path = "microsoft/phi-2"
5
- model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype="auto")
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")