diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp
index 70e837d6eb2f..8f526a93fb59 100644
--- a/src/models/lfm2.cpp
+++ b/src/models/lfm2.cpp
@@ -242,6 +242,8 @@ llama_model_lfm2::graph<iswa>::graph(const llama_model & model, const llm_graph_
     ggml_tensor * inp_out_ids = build_inp_out_ids();
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = cur;
+
         const bool is_moe_layer = il >= static_cast<int>(hparams.n_layer_dense_lead);
 
         auto * prev_cur = cur;