llamamodel: fix embedding crash for >512 tokens after #2310 (#2383)

Signed-off-by: Jared Van Bortel <[email protected]>
nomic-ai · May 29, 2024 · e94177e · e94177e
1 parent f047f38
commit e94177e
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/gpt4all-backend/llamamodel.cpp b/gpt4all-backend/llamamodel.cpp
@@ -386,7 +386,8 @@ bool LLamaModel::loadModel(const std::string &modelPath, int n_ctx, int ngl)
     bool isEmbedding = is_embedding_arch(llama_model_arch(d_ptr->model));
     const int n_ctx_train = llama_n_ctx_train(d_ptr->model);
     if (isEmbedding) {
-        d_ptr->ctx_params.n_batch = n_ctx;
+        d_ptr->ctx_params.n_batch  = n_ctx;
+        d_ptr->ctx_params.n_ubatch = n_ctx;
     } else {
         if (n_ctx > n_ctx_train) {
             std::cerr << "warning: model was trained on only " << n_ctx_train << " context tokens ("