From f10e67d22fb56d4205d7a4c8406b212f7ed42f73 Mon Sep 17 00:00:00 2001
From: Sunita Nadampalli <nadampal@amazon.com>
Date: Thu, 25 Jan 2024 21:47:55 +0000
Subject: [PATCH] ggml: softmax op: update the n_task calculation

---
 ggml.c | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ggml.c b/ggml.c
index ca98fde8ab2398..ebd24c34e82a02 100644
--- a/ggml.c
+++ b/ggml.c
@@ -16597,7 +16597,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
             } break;
         case GGML_OP_SOFT_MAX:
             {
-                n_tasks = MIN(MIN(4, n_threads), ggml_nrows(node->src[0]));
+                n_tasks = MIN(n_threads, MAX(ggml_nrows(node->src[0]) - 1, 1));
             } break;
         case GGML_OP_CONV_TRANSPOSE_1D:
             {