diff options
Diffstat (limited to 'research')
| -rw-r--r-- | research/sander/extract.go | 10 |
1 files changed, 9 insertions, 1 deletions
diff --git a/research/sander/extract.go b/research/sander/extract.go index c72fe7f..c9a2841 100644 --- a/research/sander/extract.go +++ b/research/sander/extract.go @@ -14,7 +14,15 @@ func (e *Experiment) Extract(db *sql.DB) error { var data []float32 var shape []int - if d, s, err := onnx.ExtractInitializer(e.model, "transformer.wte.weight"); err != nil { + var model *onnx.Model + + if m, err := onnx.NewModel(e.model); err != nil { + model = m + } else { + return err + } + + if d, s, err := model.ExtractInitializer("transformer.wte.weight"); err != nil { return err } else { data = d |