Rename log probs output

1 files changed, 2 insertions, 2 deletions

diff --git a/gpt2/model.go b/gpt2/model.go
index d0d1c2f..2b23cdf 100644
--- a/gpt2/model.go
+++ b/gpt2/model.go
@@ -155,7 +155,7 @@ func (m *Model) Generate(prompt []int64, steps int64, logits *[][]float32) ([]in
 
 func (m *Model) Score(tokens []int64, logProbs *[]float32) error {
 	if !m.withLogProbs {
-		panic("score requires log_probs output")
+		panic("score requires token_logprobs output")
 	}
 
 	if err := m.allocator.Init(tokens); err != nil {
@@ -167,7 +167,7 @@ func (m *Model) Score(tokens []int64, logProbs *[]float32) error {
 	}
 
 	if logProbs != nil {
-		d := m.allocator.Value("log_probs").(*ort.Tensor[float32]).GetData()
+		d := m.allocator.Value("token_logprobs").(*ort.Tensor[float32]).GetData()
 
 		*logProbs = append(*logProbs, d...)
 	}