diff options
Diffstat (limited to 'gpt2/allocator.go')
| -rw-r--r-- | gpt2/allocator.go | 8 |
1 files changed, 4 insertions, 4 deletions
diff --git a/gpt2/allocator.go b/gpt2/allocator.go index 3736b36..2609aeb 100644 --- a/gpt2/allocator.go +++ b/gpt2/allocator.go @@ -69,7 +69,7 @@ func (a *Allocator) OutputNames() []string { } if a.withLogProbs { - names = append(names, "log_probs") + names = append(names, "token_logprobs") } if a.withCache { @@ -186,7 +186,7 @@ func (a *Allocator) initOutputs(tokens []int64) error { return err } - names = append(names, "log_probs") + names = append(names, "token_logprobs") } if !a.withCache { @@ -419,11 +419,11 @@ func (a *Allocator) logits(tokens []int64, force bool) error { } func (a *Allocator) logProbs(tokens []int64, force bool) error { - const name = "log_probs" + const name = "token_logprobs" if _, ok := a.values[name]; ok { if !force { - panic("log_probs already allocated") + panic("token_logprobs already allocated") } _ = a.values[name].Destroy() |