summaryrefslogtreecommitdiff
path: root/onnx
diff options
context:
space:
mode:
Diffstat (limited to 'onnx')
-rw-r--r--onnx/cmd/extract/main.go3
1 files changed, 2 insertions, 1 deletions
diff --git a/onnx/cmd/extract/main.go b/onnx/cmd/extract/main.go
index d8f8358..32b8c5d 100644
--- a/onnx/cmd/extract/main.go
+++ b/onnx/cmd/extract/main.go
@@ -5,11 +5,12 @@ import (
"log"
"go.jknobloc.com/x/onnx"
+ "go.jknobloc.com/x/shelf"
"go.jknobloc.com/x/tensor"
)
func main() {
- data, shape, err := onnx.ExtractInitializer("gpt2/models/base/model.onnx", "transformer.wte.weight")
+ data, shape, err := onnx.ExtractInitializer(shelf.Abs("models/gpt2/model.onnx"), "transformer.wte.weight")
if err != nil {
log.Fatal(err)