gpt-2 : fix off-by-one error in batching logic

1 year ago · 99f1afb613
parent 64efeceabd
commit 99f1afb613
1 changed files with 1 additions and 1 deletions
--- a/examples/gpt-2/main.cpp
+++ b/examples/gpt-2/main.cpp
@ -756,7 +756,7 @@ int main(int argc, char ** argv) {
            // if here, it means we are still processing the input prompt
            for (int k = i; k < embd_inp.size(); k++) {
                embd.push_back(embd_inp[k]);
-                if (embd.size() > params.n_batch) {
+                if (embd.size() >= params.n_batch) {
                    break;
                }
            }