From 9c8e33df6fe94c7ef1eb4d6708348036d3fdd630 Mon Sep 17 00:00:00 2001
From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com>
Date: Mon, 5 Dec 2022 17:55:45 +0530
Subject: [PATCH] fixes

---
 src/pet/pet_model.py   | 4 ++--
 src/pet/utils/other.py | 2 --
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index f71aa04..045bb68 100644
--- a/src/pet/pet_model.py
+++ b/src/pet/pet_model.py
@@ -178,7 +178,7 @@ class PETModelForSequenceClassification(PETModel):
                 break
 
         # to make sure classifier layer is trainable
-        _set_trainable(self.base_model)
+        _set_trainable(self)
 
     def forward(
         self,
@@ -540,7 +540,7 @@ class PETModelForTokenClassification(PETModel):
                 break
 
         # to make sure classifier layer is trainable
-        _set_trainable(self.base_model)
+        _set_trainable(self)
 
     def forward(
         self,
diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py
index b7a81d5..74659d7 100644
--- a/src/pet/utils/other.py
+++ b/src/pet/utils/other.py
@@ -42,8 +42,6 @@ def _set_trainable(model):
         for name, param in model.named_parameters():
             if any(module_name in name for module_name in model.modules_to_save):
                 param.requires_grad = True
-            else:
-                param.requires_grad = False
 
 
 def fsdp_auto_wrap_policy(model):