add examples

src/pet/__init__.py

@@ -31,7 +31,7 @@ from .utils import (
     TaskType,
     bloom_model_postprocess_past_key_value,
     get_pet_model_state_dict,
+    pet_model_load_and_dispatch,
     set_pet_model_state_dict,
     shift_tokens_right,
-    pet_model_load_and_dispatch,
 )

src/pet/utils/__init__.py

@@ -4,4 +4,4 @@
 
 from .config import PETConfig, PETType, PromptLearningConfig, TaskType
 from .other import _set_trainable, bloom_model_postprocess_past_key_value, shift_tokens_right
-from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict, pet_model_load_and_dispatch
+from .save_and_load import get_pet_model_state_dict, pet_model_load_and_dispatch, set_pet_model_state_dict

src/pet/utils/save_and_load.py

@@ -88,5 +88,6 @@ def pet_model_load_and_dispatch(model, pet_model_state_dict, pet_config, max_mem
     if model.pet_config.pet_type == PETType.LORA:
         add_hook_to_module(model.base_model.model, hook)
     else:
+        remove_hook_from_submodules(model.prompt_encoder)
         add_hook_to_module(model.base_model, hook)
     return model