Browse Source

fix(runner): Set logits to 0 if false on Batch.Add

https://github.com/ollama/ollama/issues/7656
Branch: Granite3StoppingBug-7656

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
Gabe Goodhart 5 months ago
parent
commit
807ace5b1f
1 changed files with 2 additions and 0 deletions
  1. 2 0
      llama/llama.go

+ 2 - 0
llama/llama.go

@@ -384,6 +384,8 @@ func (b *Batch) Add(token int, embed []float32, pos int, logits bool, seqIds ...
 
 
 	if logits {
 	if logits {
 		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 1
 		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 1
+	} else {
+		unsafe.Slice(b.c.logits, b.allocSize())[b.c.n_tokens] = 0
 	}
 	}
 
 
 	b.c.n_tokens += 1
 	b.c.n_tokens += 1