summaryrefslogtreecommitdiff
path: root/gpt2/allocator.go
diff options
context:
space:
mode:
Diffstat (limited to 'gpt2/allocator.go')
-rw-r--r--gpt2/allocator.go12
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
}