diff --git a/src/models/mimo2.cpp b/src/models/mimo2.cpp
index b6d7acedac..2e9ebe621c 100644
--- a/src/models/mimo2.cpp
+++ b/src/models/mimo2.cpp
@@ -102,9 +102,12 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
 
     const float v_scale = hparams.f_attn_value_scale;
     const bool emit_h_nextn = cparams.embeddings_nextn;
-    const bool crop_last_layer = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked);
+    const bool extract_final_inp = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
+    const bool crop_last_layer = inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked) && !extract_final_inp;
 
     for (int il = 0; il < n_layer; ++il) {
+        res->t_layer_inp[il] = inpL;
+
         ggml_tensor * inpSA = inpL;
 
         uint32_t n_head_l    = hparams.n_head(il);
@@ -232,6 +235,14 @@ llama_model_mimo2::graph::graph(const llama_model & model, const llm_graph_param
 
     cur = inpL;
 
+    if (extract_final_inp) {
+        res->t_layer_inp[n_layer] = cur;
+
+        if (inp_out_ids && (!emit_h_nextn || cparams.embeddings_nextn_masked)) {
+            cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+        }
+    }
+
     if (emit_h_nextn) {
         cb(cur, "h_nextn", -1);
         res->t_h_nextn = cur;