From 992422100fd86adf954bd438baaa5d72ca8d0171 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Wed, 30 Nov 2022 18:27:10 +0530
Subject: [PATCH] fix

---
 src/pet/pet_model.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index 3f0f078..a3f6fa6 100644
--- a/src/pet/pet_model.py
+++ b/src/pet/pet_model.py
@@ -15,6 +15,7 @@ class PETModel(torch.nn.Module):
         super().__init__()
         self.pet_config = pet_config
         self.base_model = model
+        self.modules_to_save = None
         if pet_config.pet_type != PETType.LORA:
             self._setup_prompt_encoder()
         else:
@@ -99,6 +100,7 @@ class PETModelForSequenceClassification(PETModel):
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
         self.config = self.base_model.config
+        self.modules_to_save = ["classifier"]
 
         for name, module in self.base_model.named_children():
             if isinstance(module, torch.nn.Linear):