Update src/llama-model.cpp

Co-authored-by: Gabe Goodhart <[email protected]>

Author: ryan-mangeno

src/llama-model.cpp

@@ -8098,6 +8098,11 @@ struct llm_build_modern_bert : public llm_graph_context {
                         Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
             cb(cur, "kqv_out", il);
 
+            if (il == n_layer - 1 && inp_out_ids) {
+                cur  = ggml_get_rows(ctx0,  cur, inp_out_ids);
+                inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+            }
+
             // re-add the layer input
             cur = ggml_add(ctx0, cur, inpL);