fix gemma 2012 token inference error

Signed-off-by: intellinjun <[email protected]>
intel · Mar 22, 2024 · 73efcd2 · 73efcd2
1 parent 507131b
commit 73efcd2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/neural_speed/models/gemma/gemma_utils.cpp b/neural_speed/models/gemma/gemma_utils.cpp
@@ -83,7 +83,7 @@ void Gemma::init(const char* path_model, model_context* ctx, int n_gpu_layer_, b
   n_head = hparams.n_head;
   n_expert = hparams.n_experts;
   n_expert_used = hparams.n_experts_used;
-  scratch = gemma_mem_req(n_layer, lctx.scratch_size_ratio);
+  scratch = gemma_mem_req(n_layer, lctx.scratch_size_ratio * 1.5);
   model.scratchs = scratch;
 }