wassname/peft
1
0
Fork 0
mirror of https://github.com/wassname/peft.git synced 2026-09-09 11:28:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

fixes

Browse Source
This commit is contained in:
Sourab Mangrulkar
2022-12-05 16:02:13 +05:30
parent 5303f0c15c
commit bb63dc34f4
1 changed files with 6 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pet/utils/save_and_load.py
+6 -6
View File
@@ -12,17 +12,17 @@ def get_pet_model_state_dict(model):
"""
if model.pet_config.pet_type == PETType.LORA:
return lora_state_dict(model)
to_return = lora_state_dict(model, bias=model.pet_config.bias)
else:
to_return = {}
state_dict = model.state_dict()
prompt_embeddings = model.get_prompt_embedding_to_save()
to_return["prompt_embeddings"] = prompt_embeddings
if model.modules_to_save is not None:
for key, value in state_dict.items():
if any(module_name in key for module_name in model.modules_to_save):
to_return[key] = value
return to_return
if model.modules_to_save is not None:
for key, value in state_dict.items():
if any(module_name in key for module_name in model.modules_to_save):
to_return[key] = value
return to_return
def set_pet_model_state_dict(model, pet_model_state_dict):