diff options
Diffstat (limited to 'gpt2/allocator.go')
| -rw-r--r-- | gpt2/allocator.go | 12 |
1 files changed, 8 insertions, 4 deletions
diff --git a/gpt2/allocator.go b/gpt2/allocator.go index 46d612d..41d2ba5 100644 --- a/gpt2/allocator.go +++ b/gpt2/allocator.go @@ -19,10 +19,10 @@ type Allocator struct { func NewAllocator(cfg Config, opts Options, batchSize int) *Allocator { return &Allocator{ - config: cfg, - options: opts, - batchSize: batchSize, - values: make(map[string]ort.Value), + config: cfg, + options: opts, + batchSize: batchSize, + values: make(map[string]ort.Value), } } @@ -98,12 +98,16 @@ func (a *Allocator) Init(tokens []int64) error { if err := a.initInputs(tokens); err != nil { a.Destroy() + a.values = make(map[string]ort.Value) + return err } if err := a.initOutputs(tokens); err != nil { a.Destroy() + a.values = make(map[string]ort.Value) + return err } |
