[LLM Runtime]Fix gptneox bug (intel#671)

Signed-off-by: intellinjun <jun.lin@intel.com>
PentesterPriyanshu · Nov 14, 2023 · ef2c479 · ef2c479
1 parent 65a645f
commit ef2c479
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/intel_extension_for_transformers/llm/runtime/graph/models/gptneox/gptneox_utils.cpp b/intel_extension_for_transformers/llm/runtime/graph/models/gptneox/gptneox_utils.cpp
@@ -45,6 +45,7 @@ void model_load_internal(const std::string& fname, model_archs arch, model_conte
   std::unique_ptr<GPTNEOX> ms(new GPTNEOX());
   ms->init(fname.c_str(), lctx, n_gpu_layers, use_mmap, use_mlock, vocab_only);
   ms->load(lctx, progress_callback, progress_callback_user_data);
+  lctx.support_jblas_kv = true;
 }
 
 void GPTNEOX::init(const char* path_model, model_context& lctx, int n_gpu_layer_, bool use_mmap_, bool use_mlock_,