blob: d8f8358cdf6624e3e5d44bd1fc85d71ad70bbc95 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
|
package main
import (
"fmt"
"log"
"go.jknobloc.com/x/onnx"
"go.jknobloc.com/x/tensor"
)
func main() {
data, shape, err := onnx.ExtractInitializer("gpt2/models/base/model.onnx", "transformer.wte.weight")
if err != nil {
log.Fatal(err)
}
t := tensor.NewDense[float32](shape, data)
eot := t.Select(0, 50256).Contiguous()
fmt.Println(eot)
}
|
