diff --git a/README.md b/README.md
index 0b4fe38..8babeed 100644
--- a/README.md
+++ b/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
diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index e1805fa..a865623 100644
--- a/src/pet/pet_model.py
+++ b/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():