From 6fceb654707022833b2be0e1efe4ac85a2655a13 Mon Sep 17 00:00:00 2001 From: 51616 Date: Sun, 22 Dec 2024 17:43:00 +0000 Subject: [PATCH] model saving/loading --- hyperlora/modeling_utils.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/hyperlora/modeling_utils.py b/hyperlora/modeling_utils.py index 872ab97..eaa5e97 100644 --- a/hyperlora/modeling_utils.py +++ b/hyperlora/modeling_utils.py @@ -325,14 +325,12 @@ class ModulatedPretrainedModel(nn.Module): self.register_module("hypernet", hypernet) self.register_module("ctx_encoder", ctx_encoder) - # TODO: fix model saving/loading def state_dict(self, *args, **kwargs): state_dict = super().state_dict(*args, **kwargs) - - # remove non-trainable params + # remove non-trainable and base_model's params for name, param in self.named_parameters(): - if not param.requires_grad: - state_dict.pop(name) + if (not param.requires_grad) or name.startswith("base_model"): + state_dict.pop(name, None) return state_dict def load_state_dict(self, state_dict: dict, *args, **kwargs):