codegeex2-6b / save_model.py
Stanislas's picture
Update implementation
d8562b1
Raw
History Blame Contribute Delete
196 Bytes
from transformers import AutoModel
model = AutoModel.from_pretrained("/mnt/vepfs/qinkai/release/codegeex2-6b/", trust_remote_code=True).cuda()
model.save_pretrained("./", max_shard_size="2000MB")