Reset allocator values map on initialization error

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
 	}