From d89882a55850b0c258dee0eef1950b191bb5e307 Mon Sep 17 00:00:00 2001 From: Sourab Mangrulkar <13534540+pacman100@users.noreply.github.com> Date: Fri, 2 Dec 2022 11:24:26 +0530 Subject: [PATCH] fix --- src/pet/tuners/prompt_tuning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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