From 10657ac011b7c32f9d272d7bd8d279edd1caf550 Mon Sep 17 00:00:00 2001 From: Jonas Knobloch Date: Mon, 6 Apr 2026 22:09:04 +0200 Subject: Fix input and output names capacity --- gpt2/allocator.go | 14 ++++++++++++-- 1 file changed, 12 insertions(+), 2 deletions(-) (limited to 'gpt2/allocator.go') diff --git a/gpt2/allocator.go b/gpt2/allocator.go index 575fac5..30fb621 100644 --- a/gpt2/allocator.go +++ b/gpt2/allocator.go @@ -26,7 +26,13 @@ func NewAllocator(config Config, withCache bool, withLogProbs bool) *Allocator { } func (a *Allocator) InputNames() []string { - names := make([]string, 0, 3+2*a.config.nLayers) + capacity := 3 + + if a.withCache { + capacity += 2 * a.config.nLayers + } + + names := make([]string, 0, capacity) names = append(names, "input_ids", "position_ids", "attention_mask") @@ -40,7 +46,11 @@ func (a *Allocator) InputNames() []string { } func (a *Allocator) OutputNames() []string { - capacity := 2*a.config.nLayers + 1 + capacity := 1 + + if a.withCache { + capacity += 2 * a.config.nLayers + } names := make([]string, 0, capacity) -- cgit v1.3.1