diff --git a/setup.py b/setup.py
index 6f74e46..fcd1b39 100644
--- a/setup.py
+++ b/setup.py
@@ -39,9 +39,10 @@ setup(
         "packaging>=20.0",
         "psutil",
         "pyyaml",
-        "torch>=1.4.0",
+        "torch>=1.13.0",
         "transformers",
         "accelerate",
+        "loralib",
     ],
     extras_require=extras,
     classifiers=[
diff --git a/src/pet/mapping.py b/src/pet/mapping.py
index 0cb2210..1f2780f 100644
--- a/src/pet/mapping.py
+++ b/src/pet/mapping.py
@@ -34,6 +34,13 @@ TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING = {
 
 
 def get_pet_config(config_dict):
+    """
+    Returns a PET config object from a dictionary.
+
+    Args:
+        config_dict (:obj:`Dict[str, Any]`):
+    """
+
     return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
 
 
@@ -93,6 +100,14 @@ def _prepare_lora_config(pet_config, model_config):
 
 
 def get_pet_model(model, pet_config):
+    """
+    Returns a PET model object from a model and a config.
+
+    Args:
+        model (:obj:`transformers.PreTrainedModel`):
+        pet_config (:obj:`transformers.PETConfig`):
+    """
+
     model_config = model.config.to_dict()
     if pet_config.pet_type != PETType.LORA:
         pet_config = _prepare_prompt_learning_config(pet_config, model_config)
diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index 1398146..e1805fa 100644
--- a/src/pet/pet_model.py
+++ b/src/pet/pet_model.py
@@ -7,10 +7,30 @@ from transformers import PreTrainedModel
 from transformers.modeling_outputs import SequenceClassifierOutput
 
 from .tuners import LoRAModel, PrefixEncoder, PromptEmbedding, PromptEncoder
-from .utils import PETConfig, PETType, TaskType, shift_tokens_right
+from .utils import PETConfig, PETType, TaskType, shift_tokens_right, _set_trainable
 
 
 class PETModel(torch.nn.Module):
+    """
+    Parameter Efficient Tuning Model. Base model encompassing various PET methods.
+
+    Args:
+        model (:obj:`PreTrainedModel`): The base transformer model used for PET.
+        pet_config (:obj:`PETConfig`): The configuration of the PET model.
+
+
+    Attributes:
+        base_model (:obj:`PreTrainedModel`): The base transformer model used for PET.
+        pet_config (:obj:`PETConfig`): The configuration of the PET model.
+        modules_to_save (:obj:`list` of :obj:`str`): The list of sub-module names to save when saving the model.
+        prompt_encoder (:obj:`PromptEncoder`): The prompt encoder used for PET if `pet_config.pet_type != PETType.LORA`.
+        prompt_tokens (:obj:`torch.Tensor`): The virtual prompt tokens used for PET if `pet_config.pet_type != PETType.LORA`.
+        transformer_backbone_name (:obj:`str`): The name of the transformer backbone in the base model
+            if `pet_config.pet_type != PETType.LORA`.
+        word_embeddings (:obj:`torch.nn.Embedding`): The word embeddings of the transformer backbone
+            in the base model if `pet_config.pet_type != PETType.LORA`.
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__()
         self.pet_config = pet_config
@@ -54,6 +74,10 @@ class PETModel(torch.nn.Module):
         ).long()
 
     def get_prompt_embedding_to_save(self):
+        """
+        Returns the prompt embedding to save when saving the model.
+        Only applocable when `pet_config.pet_type != PETType.LORA`.
+        """
         prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(1, -1).to(self.base_model.device)
         if self.pet_config.pet_type == PETType.PREFIX_TUNING:
             prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
@@ -61,6 +85,10 @@ class PETModel(torch.nn.Module):
         return prompt_embeddings[0].detach().cpu()
 
     def get_prompt(self, batch_size):
+        """
+        Returns the virtual prompts to use for PET.
+        Only applocable when `pet_config.pet_type != PETType.LORA`.
+        """
         prompt_tokens = self.prompt_tokens.unsqueeze(0).expand(batch_size, -1).to(self.base_model.device)
         if self.pet_config.pet_type == PETType.PREFIX_TUNING:
             prompt_tokens = prompt_tokens[:, : self.pet_config.num_virtual_tokens]
@@ -92,6 +120,9 @@ class PETModel(torch.nn.Module):
             return prompts
 
     def print_trainable_parameters(self):
+        """
+        Prints the number of trainable parameters in the model.
+        """
         trainable_params = 0
         all_param = 0
         for _, param in self.named_parameters():
@@ -104,6 +135,41 @@ class PETModel(torch.nn.Module):
 
 
 class PETModelForSequenceClassification(PETModel):
+    """
+    PET model for sequence classification tasks.
+
+    Args:
+        model (:obj:`PreTrainedModel`): Base transformer model
+        pet_config (:obj:`PETConfig`): PET config.
+
+    Attributes:
+        config (:obj:`PretrainedConfig`): The configuration object of the base model.
+        cls_layer_name (:obj:`str`): The name of the classification layer.
+
+    Example::
+
+        >>> from transformers import AutoModelForSequenceClassification
+        >>> from pet import PETModelForSequenceClassification, get_pet_config
+        >>> config = {
+                'pet_type': 'PREFIX_TUNING',
+                'task_type': 'SEQ_CLS',
+                'inference_mode': False,
+                'num_virtual_tokens': 20,
+                'token_dim': 768,
+                'num_transformer_submodules': 1,
+                'num_attention_heads': 12,
+                'num_layers': 12,
+                'encoder_hidden_size': 768,
+                'prefix_projection': False,
+                'postprocess_past_key_value_function': None
+            }
+        >>> pet_config = get_pet_config(config)
+        >>> model = AutoModelForSequenceClassification.from_pretrained("bert-base-cased")
+        >>> pet_model = PETModelForSequenceClassification(model, pet_config)
+        >>> pet_model.print_trainable_parameters()
+        trainable params: 370178 || all params: 108680450 || trainable%: 0.3406113979101117
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
         self.config = self.base_model.config
@@ -114,6 +180,9 @@ class PETModelForSequenceClassification(PETModel):
                 self.cls_layer_name = name
                 break
 
+        # to make sure classifier layer is trainable
+        _set_trainable(self.base_model)
+
     def forward(
         self,
         input_ids=None,
@@ -160,7 +229,7 @@ class PETModelForSequenceClassification(PETModel):
         )
 
         if self.pet_config.pet_type == PETType.PREFIX_TUNING:
-            return self.prefix_tuning_forward(input_ids=input_ids, **kwargs)
+            return self._prefix_tuning_forward(input_ids=input_ids, **kwargs)
         else:
             if kwargs.get("token_type_ids", None) is not None:
                 kwargs["token_type_ids"] = torch.cat(
@@ -176,7 +245,7 @@ class PETModelForSequenceClassification(PETModel):
             inputs_embeds = torch.cat((prompts, inputs_embeds), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
-    def prefix_tuning_forward(
+    def _prefix_tuning_forward(
         self,
         input_ids=None,
         attention_mask=None,
@@ -249,9 +318,40 @@ class PETModelForSequenceClassification(PETModel):
 
 
 class PETModelForCausalLM(PETModel):
+    """
+    PET model for Causal LM
+
+    Args:
+        model (:obj:`PreTrainedModel`): Base transformer model
+        pet_config (:obj:`PETConfig`): PET config.
+
+
+    Example::
+
+        >>> from transformers import AutoModelForCausalLM
+        >>> from pet import PETModelForCausalLM, get_pet_config
+        >>> config = {
+                'pet_type': 'PREFIX_TUNING',
+                'task_type': 'CAUSAL_LM',
+                'inference_mode': False,
+                'num_virtual_tokens': 20,
+                'token_dim': 1280,
+                'num_transformer_submodules': 1,
+                'num_attention_heads': 20,
+                'num_layers': 36,
+                'encoder_hidden_size': 1280,
+                'prefix_projection': False,
+                'postprocess_past_key_value_function': None
+            }
+        >>> pet_config = get_pet_config(config)
+        >>> model = AutoModelForCausalLM.from_pretrained("gpt2-large")
+        >>> pet_model = PETModelForCausalLM(model, pet_config)
+        >>> pet_model.print_trainable_parameters()
+        trainable params: 1843200 || all params: 775873280 || trainable%: 0.23756456724479544
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.config = self.base_model.config
 
     def forward(
         self,
@@ -318,9 +418,40 @@ class PETModelForCausalLM(PETModel):
 
 
 class PETModelForSeq2SeqLM(PETModel):
+    """
+    PET model for Seq2Seq LM
+
+    Args:
+        model (:obj:`PreTrainedModel`): Base transformer model
+        pet_config (:obj:`PETConfig`): PET config.
+
+
+    Example::
+
+        >>> from transformers import AutoModelForSeq2SeqLM
+        >>> from pet import PETModelForSeq2SeqLM, get_pet_config
+        >>> config = {
+                'pet_type': 'LORA',
+                'task_type': 'SEQ_2_SEQ_LM',
+                'inference_mode': False,
+                'r': 8,
+                'target_modules': ['q', 'v'],
+                'lora_alpha': 32,
+                'lora_dropout': 0.1,
+                'merge_weights': False,
+                'fan_in_fan_out': False,
+                'enable_lora': None,
+                'bias': 'none'
+            }
+        >>> pet_config = get_pet_config(config)
+        >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
+        >>> pet_model = PETModelForSeq2SeqLM(model, pet_config)
+        >>> pet_model.print_trainable_parameters()
+        trainable params: 884736 || all params: 223843584 || trainable%: 0.3952474242013566
+    """
+
     def __init__(self, model, pet_config: PETConfig):
         super().__init__(model, pet_config)
-        self.config = self.base_model.config
 
     def forward(
         self,
diff --git a/src/pet/tuners/lora.py b/src/pet/tuners/lora.py
index eb30bad..9c3899a 100644
--- a/src/pet/tuners/lora.py
+++ b/src/pet/tuners/lora.py
@@ -13,6 +13,21 @@ from ..utils import PETConfig
 
 @dataclass
 class LoRAConfig(PETConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.LoRA`.
+
+    Args:
+        r: (:obj:`init`): LoRA attention dimension
+        target_modules (:obj: list of :obj: str): The names of the modules to apply LoRA to.
+        lora_alpha (:obj: float): The alpha parameter for LoRA scaling.
+        lora_dropout (:obj: float): The dropout probability for LoRA layers.
+        merge_weights (:obj: bool):
+            Whether to merge the weights of the LoRA layers with the base transformer model in `eval` mode.
+        fan_in_fan_out (:obj: bool): Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        enable_lora (:obj: list of :obj: bool): Used with `lora.MergedLinear`.
+        bias (:obj: str): Bias type for LoRA. Can be 'none', 'all' or 'lora_only'
+    """
+
     r: int = field(default=8, metadata={"help": "LoRA attention dimension"})
     target_modules: Optional[list] = field(default=None, metadata={"help": "List of modules to replace with LoRA"})
     lora_alpha: int = field(default=None, metadata={"help": "LoRA alpha"})
@@ -29,14 +44,44 @@ class LoRAConfig(PETConfig):
 
 
 class LoRAModel(torch.nn.Module):
+    """
+    Creates Low Rank Adapter (LoRA) model from a pretrained transformers model.
+
+    Args:
+        model (:obj:`transformers.PreTrainedModel`): The model to be adapted.
+        config (:obj:`LoRAConfig`): The configuration of the LoRA model.
+
+    Returns:
+        :obj:`torch.nn.Module`: The LoRA model.
+
+    Example::
+
+        >>> from transformers import AutoModelForSeq2SeqLM, LoRAConfig
+        >>> from pet import LoRAModel, LoRAConfig
+        >>> config = LoRAConfig(
+            pet_type="LORA",
+            task_type="SEQ_2_SEQ_LM",
+            r=8,
+            lora_alpha=32,
+            target_modules=["q", "v"],
+            lora_dropout=0.01,
+            )
+        >>> model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
+        >>> lora_model = LoRAModel(config, model)
+
+    Attributes:
+        model (:obj:`transformers.PreTrainedModel`): The model to be adapted.
+        config (:obj:`LoRAConfig`): The configuration of the LoRA model.
+    """
+
     def __init__(self, config, model):
         super().__init__()
         self.config = config
         self.model = model
-        self.find_and_replace()
+        self._find_and_replace()
         mark_only_lora_as_trainable(self.model, self.config.bias)
 
-    def find_and_replace(self):
+    def _find_and_replace(self):
         kwargs = {
             "r": self.config.r,
             "lora_alpha": self.config.lora_alpha,
@@ -47,23 +92,23 @@ class LoRAModel(torch.nn.Module):
         key_list = [key for key, _ in self.model.named_modules()]
         for key in key_list:
             if any(key.endswith(target_key) for target_key in self.config.target_modules):
-                parent, target, target_name = self.get_submodules(key)
-                # print(parent, target, target_name)
+                parent, target, target_name = self._get_submodules(key)
+                bias = target.bias is not None
                 if isinstance(target, torch.nn.Linear):
-                    new_module = lora.Linear(target.in_features, target.out_features, **kwargs)
+                    new_module = lora.Linear(target.in_features, target.out_features, bias=bias, **kwargs)
                 elif isinstance(target, Conv1D):
                     kwargs.update({"enable_lora": self.config.enable_lora})
                     in_features, out_features = target.weight.shape
-                    new_module = lora.MergedLinear(in_features, out_features, **kwargs)
-                self.replace_module(parent, target_name, new_module, target)
+                    new_module = lora.MergedLinear(in_features, out_features, bias=bias, **kwargs)
+                self._replace_module(parent, target_name, new_module, target)
 
-    def get_submodules(self, key):
+    def _get_submodules(self, key):
         parent = self.model.get_submodule(".".join(key.split(".")[:-1]))
         target_name = key.split(".")[-1]
         target = self.model.get_submodule(key)
         return parent, target, target_name
 
-    def replace_module(self, parent_module, child_name, new_module, old_module):
+    def _replace_module(self, parent_module, child_name, new_module, old_module):
         setattr(parent_module, child_name, new_module)
         new_module.weight = old_module.weight
         if old_module.bias is not None:
diff --git a/src/pet/tuners/p_tuning.py b/src/pet/tuners/p_tuning.py
index bb5a5a8..0700ba8 100644
--- a/src/pet/tuners/p_tuning.py
+++ b/src/pet/tuners/p_tuning.py
@@ -14,21 +14,33 @@ class PromptEncoderReparameterizationType(str, enum.Enum):
 
 @dataclass
 class PromptEncoderConfig(PromptLearningConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PromptEncoder`.
+
+    Args:
+        encoder_reparameterization_type
+            (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]):
+            The type of reparameterization to use.
+        encoder_hidden_size (:obj:`int`): The hidden size of the prompt encoder.
+        encoder_num_layers (:obj:`int`): The number of layers of the prompt encoder.
+        encoder_dropout (:obj:`float`): The dropout probability of the prompt encoder.
+    """
+
     encoder_reparameterization_type: Union[str, PromptEncoderReparameterizationType] = field(
         default=PromptEncoderReparameterizationType.MLP,
         metadata={"help": "How to reparameterize the prompt encoder"},
     )
     encoder_hidden_size: int = field(
         default=None,
-        metadata={"help": "The hidden size of the prompt encoder reparameterization"},
+        metadata={"help": "The hidden size of the prompt encoder"},
     )
     encoder_num_layers: int = field(
         default=2,
-        metadata={"help": "The number of layers of the prompt encoder reparameterization"},
+        metadata={"help": "The number of layers of the prompt encoder"},
     )
     encoder_dropout: float = field(
         default=0.0,
-        metadata={"help": "The dropout of the prompt encoder reparameterization"},
+        metadata={"help": "The dropout of the prompt encoder"},
     )
 
 
@@ -37,6 +49,43 @@ class PromptEncoderConfig(PromptLearningConfig):
 class PromptEncoder(torch.nn.Module):
     """
     The prompt encoder network that is used to generate the virtual token embeddings for p-tuning.
+
+    Args:
+        config (:class:`PromptEncoderConfig`): The configuration of the prompt encoder.
+
+    Example::
+
+        >>> from pet import PromptEncoder, PromptEncoderConfig
+        >>> config = PromptEncoderConfig(
+                pet_type="P_TUNING",
+                task_type="SEQ_2_SEQ_LM",
+                num_virtual_tokens=20,
+                token_dim=768,
+                num_transformer_submodules=1,
+                num_attention_heads=12,
+                num_layers=12,
+                encoder_reparameterization_type="MLP",
+                encoder_hidden_size=768
+            )
+        >>> prompt_encoder = PromptEncoder(config)
+
+    Attributes:
+        embedding (:class:`~torch.nn.Embedding`): The embedding layer of the prompt encoder.
+        mlp_head (:class:`~torch.nn.Sequential`): The MLP head of the prompt encoder if `inference_mode=False`.
+        lstm_head (:class:`~torch.nn.LSTM`):
+            The LSTM head of the prompt encoder if `inference_mode=False` and `encoder_reparameterization_type="LSTM"`.
+        token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model.
+        input_size (:obj:`int`): The input size of the prompt encoder.
+        output_size (:obj:`int`): The output size of the prompt encoder.
+        hidden_size (:obj:`int`): The hidden size of the prompt encoder.
+        total_virtual_tokens (:obj:`int`): The total number of virtual tokens of the prompt encoder.
+        encoder_type (:obj:Union[:class:`PromptEncoderReparameterizationType`, :obj:`str`]):
+            The encoder type of the prompt encoder.
+
+
+    Input shape: (batch_size, total_virtual_tokens)
+
+    Output shape: (batch_size, total_virtual_tokens, token_dim)
     """
 
     def __init__(self, config):
diff --git a/src/pet/tuners/prefix_tuning.py b/src/pet/tuners/prefix_tuning.py
index d709baf..cb4c575 100644
--- a/src/pet/tuners/prefix_tuning.py
+++ b/src/pet/tuners/prefix_tuning.py
@@ -8,6 +8,15 @@ from ..utils import PromptLearningConfig
 
 @dataclass
 class PrefixTuningConfig(PromptLearningConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PrefixEncoder`.
+
+    Args:
+        encoder_hidden_size (:obj: int): The hidden size of the prompt encoder.
+        prefix_projection (:obj: bool): Whether to project the prefix embeddings.
+        postprocess_past_key_value_function (:obj: Optional[Callable]): The function to postprocess the past key value.
+    """
+
     encoder_hidden_size: int = field(
         default=None,
         metadata={"help": "The hidden size of the encoder"},
@@ -28,6 +37,31 @@ class PrefixEncoder(torch.nn.Module):
     r"""
     The torch.nn model to encode the prefix
 
+    Args:
+        config (:class:`PrefixTuningConfig`): The configuration of the prefix encoder.
+
+    Example::
+
+        >>> from pet import PrefixEncoder, PrefixTuningConfig
+        >>> config = PrefixTuningConfig(
+                pet_type="PREFIX_TUNING",
+                task_type="SEQ_2_SEQ_LM",
+                num_virtual_tokens=20,
+                token_dim=768,
+                num_transformer_submodules=1,
+                num_attention_heads=12,
+                num_layers=12,
+                encoder_hidden_size=768
+            )
+        >>> prefix_encoder = PrefixEncoder(config)
+
+
+    Attributes:
+        embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prefix encoder.
+        trans (:obj:`torch.nn.Sequential`): The two-layer MLP to transform the prefix embeddings
+            if :obj:`prefix_projection` is :obj:`True`.
+        prefix_projection (:obj:`bool`): Whether to project the prefix embeddings.
+
     Input shape: (batch_size, num_virtual_tokens)
 
     Output shape: (batch_size, num_virtual_tokens, 2*layers*hidden)
diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py
index fdad421..5ee6ef1 100644
--- a/src/pet/tuners/prompt_tuning.py
+++ b/src/pet/tuners/prompt_tuning.py
@@ -15,6 +15,17 @@ class PromptTuningInit(str, enum.Enum):
 
 @dataclass
 class PromptTuningConfig(PromptLearningConfig):
+    """
+    This is the configuration class to store the configuration of a :class:`~pet.PromptEmbedding`.
+
+    Args:
+        prompt_tuning_init (:obj:Union[:class:`PromptTuningInit`, :obj:`str`]): The initialization of the prompt embedding.
+        prompt_tuning_init_text (:obj: Optional[:obj:`str`]): The text to initialize the prompt embedding.
+            Only used if `prompt_tuning_init` is `TEXT`
+        tokenizer_name_or_path (:obj: Optional[:obj:`str`]): The name or path of the tokenizer.
+            Only used if `prompt_tuning_init` is `TEXT`
+    """
+
     prompt_tuning_init: Union[PromptTuningInit, str] = field(
         default=PromptTuningInit.RANDOM,
         metadata={"help": "How to initialize the prompt tuning parameters"},
@@ -34,6 +45,40 @@ class PromptTuningConfig(PromptLearningConfig):
 
 
 class PromptEmbedding(torch.nn.Module):
+    """
+    The model to encode virtual tokens into prompt embeddings.
+
+    Args:
+        config (:class:`PromptTuningConfig`): The configuration of the prompt embedding.
+        word_embeddings (:obj:`torch.nn.Module`): The word embeddings of the base transformer model.
+
+    Attributes:
+        embedding (:obj:`torch.nn.Embedding`): The embedding layer of the prompt embedding.
+
+    Example::
+
+        >>> from pet import PromptEmbedding, PromptTuningConfig
+        >>> config = PromptTuningConfig(
+                pet_type="PROMPT_TUNING",
+                task_type="SEQ_2_SEQ_LM",
+                num_virtual_tokens=20,
+                token_dim=768,
+                num_transformer_submodules=1,
+                num_attention_heads=12,
+                num_layers=12,
+                prompt_tuning_init="TEXT",
+                prompt_tuning_init_text="Predict if sentiment of this review is positive, negative or neutral",
+                tokenizer_name_or_path="t5-base",
+            )
+        >>> # t5_model.shared is the word embeddings of the base model
+        >>> prompt_embedding = PromptEmbedding(config, t5_model.shared)
+
+
+    Input Shape: (batch_size, total_virtual_tokens)
+
+    Output Shape: (batch_size, total_virtual_tokens, token_dim)
+    """
+
     def __init__(self, config, word_embeddings):
         super().__init__()
 
@@ -42,9 +87,9 @@ class PromptEmbedding(torch.nn.Module):
         if config.prompt_tuning_init == PromptTuningInit.TEXT:
             from transformers import AutoTokenizer
 
-            self.tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
-            self.init_text = config.prompt_tuning_init_text
-            init_token_ids = self.tokenizer(self.init_text)["input_ids"]
+            tokenizer = AutoTokenizer.from_pretrained(config.tokenizer_name_or_path)
+            init_text = config.prompt_tuning_init_text
+            init_token_ids = tokenizer(init_text)["input_ids"]
             # Trim or iterate until num_text_tokens matches total_virtual_tokens
             num_text_tokens = len(init_token_ids)
             if num_text_tokens > total_virtual_tokens:
diff --git a/src/pet/utils/__init__.py b/src/pet/utils/__init__.py
index aa359db..582dad4 100644
--- a/src/pet/utils/__init__.py
+++ b/src/pet/utils/__init__.py
@@ -3,5 +3,5 @@
 # module, but to preserve other warnings. So, don't check this module at all
 
 from .config import PETConfig, PETType, PromptLearningConfig, TaskType
-from .other import bloom_model_postprocess_past_key_value, shift_tokens_right
+from .other import bloom_model_postprocess_past_key_value, shift_tokens_right, _set_trainable
 from .save_and_load import get_pet_model_state_dict, set_pet_model_state_dict
diff --git a/src/pet/utils/config.py b/src/pet/utils/config.py
index d989601..984016b 100644
--- a/src/pet/utils/config.py
+++ b/src/pet/utils/config.py
@@ -19,7 +19,12 @@ class TaskType(str, enum.Enum):
 @dataclass
 class PETConfig:
     """
-    This is the configuration class to store the configuration of a :class:`~pet.PETModel`.
+    This is the base configuration class to store the configuration of a :class:`~pet.PETModel`.
+
+    Args:
+        pet_type (:obj:Union[:class:`~pet.utils.config.PETType`, :obj:`str`]): The type of PET method to use.
+        task_type (:obj:Union[:class:`~pet.utils.config.TaskType`, :obj:`str`]): The type of task to perform.
+        inference_mode (:obj:`bool`, defaults to :obj:`False`): Whether to use the PET model in inference mode.
     """
 
     pet_type: Union[str, PETType] = field(default=None, metadata={"help": "PET type"})
@@ -29,8 +34,21 @@ class PETConfig:
 
 @dataclass
 class PromptLearningConfig(PETConfig):
+    """
+    This is the base configuration class to store the configuration of a :obj:Union[:class:`~pet.PrefixTuning`, :class:`~pet.PromptEncoder`, :class:`~pet.PromptTuning`].
+
+    Args:
+        num_virtual_tokens (:obj:`int`): The number of virtual tokens to use.
+        token_dim (:obj:`int`): The hidden embedding dimension of the base transformer model.
+        num_transformer_submodules (:obj:`int`): The number of transformer submodules in the base transformer model.
+        num_attention_heads (:obj:`int`): The number of attention heads in the base transformer model.
+        num_layers (:obj:`int`): The number of layers in the base transformer model.
+    """
+
     num_virtual_tokens: int = field(default=None, metadata={"help": "Number of virtual tokens"})
-    token_dim: int = field(default=None, metadata={"help": "Dimension of virtual tokens"})
+    token_dim: int = field(
+        default=None, metadata={"help": "The hidden embedding dimension of the base transformer model"}
+    )
     num_transformer_submodules: Optional[int] = field(default=1, metadata={"help": "Number of transformer submodules"})
     num_attention_heads: Optional[int] = field(default=None, metadata={"help": "Number of attention heads"})
     num_layers: Optional[int] = field(default=None, metadata={"help": "Number of transformer layers"})
diff --git a/src/pet/utils/other.py b/src/pet/utils/other.py
index ac8b329..c6274db 100644
--- a/src/pet/utils/other.py
+++ b/src/pet/utils/other.py
@@ -19,6 +19,11 @@ def bloom_model_postprocess_past_key_value(past_key_values):
 def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
     """
     Shift input ids one token to the right.
+
+    Args:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`): input ids
+        pad_token_id (:obj:`int`): The id of the `padding` token.
+        decoder_start_token_id (:obj:`int`): The id of the `start` token.
     """
     shifted_input_ids = input_ids.new_zeros(input_ids.shape)
     shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
@@ -30,3 +35,12 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start
     shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
 
     return shifted_input_ids
+
+
+def _set_trainable(model):
+    if model.modules_to_save is not None:
+        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
diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py
index e2eea05..1418dfa 100644
--- a/src/pet/utils/save_and_load.py
+++ b/src/pet/utils/save_and_load.py
@@ -4,6 +4,13 @@ from .config import PETType
 
 
 def get_pet_model_state_dict(model):
+    """
+    Get the state dict of the PET model.
+
+    Args:
+        model (:obj:`PETModel`): The PET model.
+    """
+
     if model.pet_config.pet_type == PETType.LORA:
         return lora_state_dict(model)
     else:
@@ -19,6 +26,14 @@ def get_pet_model_state_dict(model):
 
 
 def set_pet_model_state_dict(model, pet_model_state_dict):
+    """
+    Set the state dict of the PET model.
+
+    Args:
+        model (:obj:`PETModel`): The PET model.
+        pet_model_state_dict (:obj:`dict`): The state dict of the PET model.
+    """
+
     model.load_state_dict(pet_model_state_dict, strict=False)
     if model.pet_config.pet_type != PETType.LORA:
         model.prompt_encoder.embedding.load_state_dict(