fix

README.md

@@ -28,6 +28,7 @@ pet_config = get_pet_config(config)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name_or_path)
 model = get_pet_model(model, pet_config)
 model.print_trainable_parameters()
+# output: 
 ```
 
 ## Models support matrix

src/pet/pet_model.py

@@ -35,6 +35,7 @@ class PETModel(torch.nn.Module):
         super().__init__()
         self.pet_config = pet_config
         self.base_model = model
+        self.config = self.base_model.config
         self.modules_to_save = None
         if pet_config.pet_type != PETType.LORA:
             self._setup_prompt_encoder()
@@ -172,7 +173,6 @@ 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():