ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Update save.py

Browse source
This commit is contained in:
Daniel Han-Chen 2024-01-27 04:20:58 +11:00
commit baeea64917
1 changed files with 1 additions and 1 deletions
Download patch file Download diff file Expand all files Collapse all files

2
unsloth/save.py
View file

@ -351,7 +351,7 @@ def unsloth_save_model(
# HF also uses a OrderedDict
from collections import OrderedDict
state_dict = OrderedDict()
state_dict["model.embed_tokens.weight"] = internal_model.embed_tokens.weight.data
state_dict["model.embed_tokens.weight"] = internal_model.model.embed_tokens.weight.data
max_vram = int(torch.cuda.get_device_properties(0).total_memory * maximum_memory_usage)