diff options
Diffstat (limited to 'onnx/cmd/extract/main.go')
| -rw-r--r-- | onnx/cmd/extract/main.go | 10 |
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) |