From e16279ed0e7854c5a0792c6600d4bdfafa0933e0 Mon Sep 17 00:00:00 2001
From: Minsoo Cheong <icycle0409@snu.ac.kr>
Date: Sat, 23 Mar 2024 21:41:22 +0900
Subject: [PATCH] assign n_batch value to n_ubatch

---
 examples/retrieval/retrieval.cpp | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/examples/retrieval/retrieval.cpp b/examples/retrieval/retrieval.cpp
index 2fa9cfa3f..467e2f027 100644
--- a/examples/retrieval/retrieval.cpp
+++ b/examples/retrieval/retrieval.cpp
@@ -108,6 +108,9 @@ int main(int argc, char ** argv) {
         return 1;
     }
 
+    // For BERT models, batch size must be equal to ubatch size
+    params.n_ubatch = params.n_batch;
+
     if (params.chunk_size <= 0) {
         fprintf(stderr, "chunk_size must be positive\n");
         return 1;