diff --git a/gpttype_adapter.cpp b/gpttype_adapter.cpp
index 1d33a2243..ba1f211bd 100644
--- a/gpttype_adapter.cpp
+++ b/gpttype_adapter.cpp
@@ -412,7 +412,7 @@ generation_outputs gpttype_generate(const generation_inputs inputs, generation_o
     {
         n_vocab = gpt2_ctx_v2.hparams.n_vocab;
     }
-    else if(file_format == FileFormat::NEOX_1)
+    else if(file_format == FileFormat::NEOX_1 || file_format == FileFormat::NEOX_2)
     {
         n_vocab = neox_ctx.hparams.n_vocab;
     }