Fix a bug about head gathering.

cryscan · Sep 7, 2023 · e61023f · e61023f
1 parent 9f6a0e2
commit e61023f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/examples/batch.rs b/examples/batch.rs
@@ -146,9 +146,9 @@ async fn run(cli: Cli) -> Result<()> {
     // [`BackedState::repeat`] is helpful if you want to create batch of states from the same input.
     let state = ModelState::new(&context, model.info(), tokens.len());
 
-    let mut num_tokens = [100usize, 200, 300, 400]
+    let mut num_tokens = [100usize, 400, 200, 300]
         .to_vec()
-        .repeat((cli.batch + prompts.len() - 1) / prompts.len())[..cli.batch]
+        .repeat((cli.batch + 3) / 4)[..cli.batch]
         .to_vec();
     loop {
         #[cfg(not(debug_assertions))]

diff --git a/src/model.rs b/src/model.rs
@@ -871,7 +871,7 @@ impl<'a> Model<'a> {
         let stack = self.request_stack(num_batch);
 
         // gather and group copy operations
-        let (head_ops, head_x) = if num_token == 1 || num_token == max_batch {
+        let (head_ops, head_x) = if num_token == 1 || num_token == num_header {
             (vec![], &buffer.ffn_x)
         } else {
             let mut start = 0;