diff --git a/src/pet/tuners/prompt_tuning.py b/src/pet/tuners/prompt_tuning.py
index 20e3a2e..5b6c313 100644
--- a/src/pet/tuners/prompt_tuning.py
+++ b/src/pet/tuners/prompt_tuning.py
@@ -77,7 +77,7 @@ class PromptEmbedding(torch.nn.Module):
         super().__init__()
 
         total_virtual_tokens = config.num_virtual_tokens * config.num_transformer_submodules
-        self.embedding = torch.nn.Embedding(total_virtual_tokens, config["token_dim"])
+        self.embedding = torch.nn.Embedding(total_virtual_tokens, config.token_dim)
         if config.prompt_tuning_init == PromptTuningInit.TEXT:
             from transformers import AutoTokenizer