diff --git a/src/pet/mapping.py b/src/pet/mapping.py
index f2a3969..e0b934a 100644
--- a/src/pet/mapping.py
+++ b/src/pet/mapping.py
@@ -37,38 +37,38 @@ def get_pet_config(config_dict):
     return PET_TYPE_TO_CONFIG_MAPPING[config_dict["pet_type"]](**config_dict)
 
 
-def _prepare_prompt_learning_config(pet_config, config):
+def _prepare_prompt_learning_config(pet_config, model_config):
     if pet_config.num_layers is None:
-        if "num_hidden_layers" in config:
-            num_layers = config["num_hidden_layers"]
-        elif "num_layers" in config:
-            num_layers = config["num_layers"]
-        elif "n_layer" in config:
-            num_layers = config["n_layer"]
+        if "num_hidden_layers" in model_config:
+            num_layers = model_config["num_hidden_layers"]
+        elif "num_layers" in model_config:
+            num_layers = model_config["num_layers"]
+        elif "n_layer" in model_config:
+            num_layers = model_config["n_layer"]
         else:
             raise ValueError("Please specify `num_layers` in `pet_config`")
         pet_config.num_layers = num_layers
 
     if pet_config.token_dim is None:
-        if "hidden_size" in config:
-            token_dim = config["hidden_size"]
-        elif "n_embd" in config:
-            token_dim = config["n_embd"]
-        elif "d_model" in config:
-            token_dim = config["d_model"]
+        if "hidden_size" in model_config:
+            token_dim = model_config["hidden_size"]
+        elif "n_embd" in model_config:
+            token_dim = model_config["n_embd"]
+        elif "d_model" in model_config:
+            token_dim = model_config["d_model"]
         else:
             raise ValueError("Please specify `token_dim` in `pet_config`")
         pet_config.token_dim = token_dim
 
     if pet_config.num_attention_heads is None:
-        if "num_attention_heads" in config:
-            num_attention_heads = config["num_attention_heads"]
-        elif "n_head" in config:
-            num_attention_heads = config["n_head"]
-        elif "num_heads" in config:
-            num_attention_heads = config["num_heads"]
-        elif "encoder_attention_heads" in config:
-            num_attention_heads = config["encoder_attention_heads"]
+        if "num_attention_heads" in model_config:
+            num_attention_heads = model_config["num_attention_heads"]
+        elif "n_head" in model_config:
+            num_attention_heads = model_config["n_head"]
+        elif "num_heads" in model_config:
+            num_attention_heads = model_config["num_heads"]
+        elif "encoder_attention_heads" in model_config:
+            num_attention_heads = model_config["encoder_attention_heads"]
         else:
             raise ValueError("Please specify `num_attention_heads` in `pet_config`")
         pet_config.num_attention_heads = num_attention_heads
@@ -79,11 +79,11 @@ def _prepare_prompt_learning_config(pet_config, config):
     return pet_config
 
 
-def _prepare_lora_config(pet_config, config):
+def _prepare_lora_config(pet_config, model_config):
     if pet_config.target_modules is None:
-        if config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
+        if model_config.model_type not in TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING:
             raise ValueError("Please specify `target_modules` in `pet_config`")
-        pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[config.model_type]
+        pet_config.target_modules = TRANSFORMERS_MODELS_TO_LORA_TARGET_MODULES_MAPPING[model_config.model_type]
     if len(pet_config.target_modules) == 1:
         pet_config.fan_in_fan_out = True
         pet_config.enable_lora = [True, False, True]
@@ -93,10 +93,10 @@ def _prepare_lora_config(pet_config, config):
 
 
 def get_pet_model(model, pet_config):
-    config = model.config.to_dict()
+    model_config = model.config.to_dict()
     if pet_config.pet_type != PETType.LORA:
-        pet_config = _prepare_prompt_learning_config(pet_config, config)
+        pet_config = _prepare_prompt_learning_config(pet_config, model_config)
     else:
-        pet_config = _prepare_lora_config(pet_config, config)
+        pet_config = _prepare_lora_config(pet_config, model_config)
 
     return MODEL_TYPE_TO_PET_MODEL_MAPPING[pet_config.task_type](model, pet_config)