Refactor ONNX extraction

1 files changed, 9 insertions, 1 deletions

diff --git a/onnx/cmd/extract/main.go b/onnx/cmd/extract/main.go
index 32b8c5d..a2c281d 100644
--- a/onnx/cmd/extract/main.go
+++ b/onnx/cmd/extract/main.go
@@ -10,7 +10,15 @@ import (
 )
 
 func main() {
-	data, shape, err := onnx.ExtractInitializer(shelf.Abs("models/gpt2/model.onnx"), "transformer.wte.weight")
+	var model *onnx.Model
+
+	if m, err := onnx.NewModel(shelf.Abs("models/gpt2/model.onnx")); err != nil {
+		log.Fatal(err)
+	} else {
+		model = m
+	}
+
+	data, shape, err := model.ExtractInitializer("transformer.wte.weight")
 
 	if err != nil {
 		log.Fatal(err)