summaryrefslogtreecommitdiff
path: root/gpt2/model.go
diff options
context:
space:
mode:
Diffstat (limited to 'gpt2/model.go')
-rw-r--r--gpt2/model.go24
1 files changed, 13 insertions, 11 deletions
diff --git a/gpt2/model.go b/gpt2/model.go
index 20dcb71..5fef956 100644
--- a/gpt2/model.go
+++ b/gpt2/model.go
@@ -12,20 +12,20 @@ import (
)
type Model struct {
- name string
- deviceID string
- config Config
- options Options
- session *ort.DynamicAdvancedSession
- allocator *Allocator
+ name string
+ deviceID string
+ config Config
+ options Options
+ session *ort.DynamicAdvancedSession
+ allocator *Allocator
}
func NewModel(name string, deviceID string, cfg Config, opts Options) *Model {
return &Model{
- name: name,
- deviceID: deviceID,
- config: cfg,
- options: opts,
+ name: name,
+ deviceID: deviceID,
+ config: cfg,
+ options: opts,
}
}
@@ -89,8 +89,10 @@ func (m *Model) Init() error {
return nil
}
-func (m *Model) Destroy() {
+func (m *Model) Destroy() error {
m.allocator.Destroy()
+
+ return m.session.Destroy()
}
func (m *Model) Generate(prompt []int64, steps int64, logits *[][]float32) ([]int64, error) {