ptuning inference table bug fix (#7015)

* remove hardcoded input and output Signed-off-by: arendu <[email protected]> * fix inf table Signed-off-by: arendu <[email protected]> --------- Signed-off-by: arendu <[email protected]> Signed-off-by: Adi Renduchintala <[email protected]>
NVIDIA · Jul 12, 2023 · a69f0f4 · a69f0f4
1 parent 40c1ccc
commit a69f0f4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py b/nemo/collections/nlp/modules/common/megatron/adapters/parallel_adapters.py
@@ -268,7 +268,7 @@ def __init__(
         # (@adithyare) the persistent=False will not pollute the indices into the state_dict of this module.
         self.register_buffer("indices", torch.LongTensor(list(range(self.virtual_tokens))), persistent=False)
         self.embedding = torch.nn.Embedding(self.virtual_tokens, self.embedding_dim)
-        self.inference_table = InferenceTable("taskname", self.embedding_dim, self.virtual_tokens)
+        self.inference_table = InferenceTable("taskname", self.output_dim, self.virtual_tokens)
         self.first = ColumnParallelLinear(
             self.embedding_dim,
             self.bottleneck_dim,