diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 3fdead4422..02bd3d10b7 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -579,10 +579,7 @@ void llama_context::resolve_fused_ops(const llama_memory_context_i * mctx, uint3
     }
 }
 
-// count the number of tensors in the graph that are flagged as inputs
-// this includes the graph nodes themselves as well as their src tensors
 static int llama_graph_n_input_tensors(ggml_cgraph * gf) {
-    // map each input tensor to the graph nodes that use it
     std::unordered_map<const ggml_tensor *, std::vector<ggml_tensor *>> users;
     for (int i = 0; i < ggml_graph_n_nodes(gf); ++i) {
         ggml_tensor * node = ggml_graph_node(gf, i);
@@ -687,7 +684,7 @@ void llama_context::sched_reserve() {
         n_splits_pp        = ggml_backend_sched_get_n_splits(sched.get());
         n_nodes_pp         = ggml_graph_n_nodes(gf);
         n_inputs_pp        = get_gf_res_reserve()->inputs.size();
-        n_input_tensors_pp = llama_graph_n_input_tensors(gf);
+        n_input_tensors_pp = this->n_intput_tensors;
     }
 
     // reserve with tg (token generation) graph to get the number of splits and nodes
@@ -700,7 +697,7 @@ void llama_context::sched_reserve() {
         n_splits_tg        = ggml_backend_sched_get_n_splits(sched.get());
         n_nodes_tg         = ggml_graph_n_nodes(gf);
         n_inputs_tg        = get_gf_res_reserve()->inputs.size();
-        n_input_tensors_tg = llama_graph_n_input_tensors(gf);
+        n_input_tensors_tg = this->n_intput_tensors;
     }
 
     // reserve again with pp graph to avoid ggml-alloc reallocations during inference
@@ -2481,6 +2478,7 @@ ggml_cgraph * llama_context::graph_reserve(
 
     auto * gf = model.build_graph(gparams);
 
+    this->n_intput_tensors = llama_graph_n_input_tensors(gf);
     this->n_outputs = save_n_outputs;
 
     // initialize scheduler with the specified graph
diff --git a/src/llama-context.h b/src/llama-context.h
index bf91daa8b5..2e69a9d5db 100644
--- a/src/llama-context.h
+++ b/src/llama-context.h
@@ -330,6 +330,7 @@ private:
     // reuse the batch_allocr to avoid unnecessary memory allocations
     std::unique_ptr<llama_batch_allocr> balloc;
 
+    uint32_t n_intput_tensors = 0; // number of tensors marked as input during the last graph reserve
     uint32_t n_outputs = 0; // number of actually-used outputs in the current ubatch or last logical batch
 
     std::vector<int32_t> output_ids; // map batch token positions to ids of the logits and embd buffers