diff --git a/src/pet/pet_model.py b/src/pet/pet_model.py
index 7124d4b..07f9669 100644
--- a/src/pet/pet_model.py
+++ b/src/pet/pet_model.py
@@ -237,7 +237,7 @@ class PETModelForSequenceClassification(PETModel):
             if inputs_embeds is None:
                 inputs_embeds = self.word_embeddings(input_ids)
             prompts = self.get_prompt(batch_size=batch_size)
-            inputs_embeds = torch.cat((inputs_embeds[:, 0, :], prompts, inputs_embeds[:, 1:, :]), dim=1)
+            inputs_embeds = torch.cat((inputs_embeds[:, 0, :].unsqueeze(1), prompts, inputs_embeds[:, 1:, :]), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
     def _prefix_tuning_forward(
@@ -680,7 +680,7 @@ class PETModelForTokenClassification(PETModel):
             if inputs_embeds is None:
                 inputs_embeds = self.word_embeddings(input_ids)
             prompts = self.get_prompt(batch_size=batch_size)
-            inputs_embeds = torch.cat((inputs_embeds[:, 0, :], prompts, inputs_embeds[:, 1:, :]), dim=1)
+            inputs_embeds = torch.cat((inputs_embeds[:, 0, :].unsqueeze(1), prompts, inputs_embeds[:, 1:, :]), dim=1)
             return self.base_model(inputs_embeds=inputs_embeds, **kwargs)
 
     def _prefix_tuning_forward(
diff --git a/src/pet/utils/save_and_load.py b/src/pet/utils/save_and_load.py
index 7dd19f7..707722f 100644
--- a/src/pet/utils/save_and_load.py
+++ b/src/pet/utils/save_and_load.py
@@ -1,7 +1,6 @@
 from loralib import lora_state_dict
 
 from .config import PETType
-from ..mapping import get_pet_model
 
 
 def get_pet_model_state_dict(model, state_dict=None):
@@ -61,6 +60,7 @@ def pet_model_load_and_dispatch(model, pet_model_state_dict, pet_config, max_mem
     """
     from accelerate import infer_auto_device_map, dispatch_model
     from accelerate.hooks import remove_hook_from_submodules, AlignDevicesHook, add_hook_to_module
+    from ..mapping import get_pet_model
 
     remove_hook_from_submodules(model)
     model = get_pet_model(model, pet_config)