From b706d682ad2fc4282425a63de1d1d7d4dd30f40c Mon Sep 17 00:00:00 2001
From: leejet <leejet714@gmail.com>
Date: Mon, 18 May 2026 23:47:10 +0800
Subject: [PATCH] fix: restore singleton dims for LLM outputs (#1518)

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

diff --git a/src/llm.hpp b/src/llm.hpp
index ab673b20..cec8b1dc 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) {