minor fix for param. names (#1495)

k2-fsa · Feb 20, 2024 · 027302c · 027302c
1 parent e59fa38
commit 027302c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/icefall/lm_wrapper.py b/icefall/lm_wrapper.py
@@ -159,7 +159,7 @@ def get_lm(self, lm_type: str, device, params: AttributeDict) -> torch.nn.Module
         """
         if lm_type == "rnn":
             model = RnnLmModel(
-                vocab_size=params.vocab_size,
+                vocab_size=params.lm_vocab_size,
                 embedding_dim=params.rnn_lm_embedding_dim,
                 hidden_dim=params.rnn_lm_hidden_dim,
                 num_layers=params.rnn_lm_num_layers,
@@ -183,7 +183,7 @@ def get_lm(self, lm_type: str, device, params: AttributeDict) -> torch.nn.Module
 
         elif lm_type == "transformer":
             model = TransformerLM(
-                vocab_size=params.vocab_size,
+                vocab_size=params.lm_vocab_size,
                 d_model=params.transformer_lm_encoder_dim,
                 embedding_dim=params.transformer_lm_embedding_dim,
                 dim_feedforward=params.transformer_lm_dim_feedforward,