From 727f180ec62f3dfec18fb6cc0dada66230afa907 Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Mon, 18 May 2026 23:28:29 +0800
Subject: [PATCH] fix: restore singleton dims for LLM outputs

---
 src/llm.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/llm.hpp b/src/llm.hpp
index ab673b208..cec8b1dcc 100644
--- a/src/llm.hpp
+++ b/src/llm.hpp
@@ -1403,7 +1403,8 @@ namespace LLM {
                                    out_layers,
                                    return_all_hidden_states);
             };
-            return take_or_empty(GGMLRunner::compute<float>(get_graph, n_threads, true));
+            return restore_trailing_singleton_dims(GGMLRunner::compute<float>(get_graph, n_threads, true),
+                                                   input_ids.dim() + 1);
         }
 
         int64_t get_num_image_tokens(int64_t t, int64_t h, int64_t w) {