diff options
Diffstat (limited to 'baselines/peft_baseline.py')
| -rw-r--r-- | baselines/peft_baseline.py | 13 |
1 files changed, 4 insertions, 9 deletions
diff --git a/baselines/peft_baseline.py b/baselines/peft_baseline.py index af1b71c..9d3b03a 100644 --- a/baselines/peft_baseline.py +++ b/baselines/peft_baseline.py @@ -87,6 +87,8 @@ class PEFTBaseline: def __init__(self, wrapper, peft_config): self.wrapper = wrapper self.device = wrapper.device + # Save reference to original model BEFORE wrapping + self._original_model = wrapper.model self.peft_model = get_peft_model(wrapper.model, peft_config) self.n_params = sum(p.numel() for p in self.peft_model.parameters() if p.requires_grad) @@ -214,14 +216,7 @@ class PEFTBaseline: def cleanup(self): """Remove adapter and restore wrapper.model to the original base model.""" - try: - base_model = self.peft_model.unload() - self.wrapper.model = base_model - except (AttributeError, NotImplementedError): - # unload() not supported for PromptTuning/PrefixTuning - # Access base model directly - self.wrapper.model = self.peft_model.base_model - if hasattr(self.wrapper.model, 'model'): - self.wrapper.model = self.wrapper.model.model + # Always restore the saved original model reference — safe for all PEFT types + self.wrapper.model = self._original_model del self.peft_model torch.cuda.empty_cache() |