minor fix

Signed-off-by: HuiyingLi <[email protected]>
HuiyingLi · Oct 23, 2024 · 903e64c · 903e64c
1 parent 795e7d6
commit 903e64c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/nemo/collections/llm/inference/base.py b/nemo/collections/llm/inference/base.py
@@ -82,7 +82,7 @@ def setup_model_and_tokenizer(
     inference_wrapped_model = GPTInferenceWrapper(
         mcore_model,
         InferenceWrapperConfig(
-            hidden_size=model.config.hidden_size,
+            hidden_size=mcore_model.config.hidden_size,
             params_dtype=params_dtype,
             inference_batch_times_seqlen_threshold=inference_batch_times_seqlen_threshold,
             padded_vocab_size=model.tokenizer.vocab_size,