summaryrefslogtreecommitdiff
path: root/gpt2/model_test.go
diff options
context:
space:
mode:
authorJonas Knobloch <jonas.knobloch@t-online.de>2026-04-08 22:24:10 +0200
committerJonas Knobloch <jonas.knobloch@t-online.de>2026-04-08 23:46:10 +0200
commit4386784c1006321c5cf0e1e005b19e0a7d7563ae (patch)
tree3d22360aaac8e4d64ca90e39b7247942073582d1 /gpt2/model_test.go
parenta006507a24b93a6af7e3cff8833fb96b48b159d7 (diff)
Refactor model options
* Add options struct
Diffstat (limited to 'gpt2/model_test.go')
-rw-r--r--gpt2/model_test.go8
1 files changed, 7 insertions, 1 deletions
diff --git a/gpt2/model_test.go b/gpt2/model_test.go
index a4de628..e562e0d 100644
--- a/gpt2/model_test.go
+++ b/gpt2/model_test.go
@@ -40,7 +40,13 @@ func fromModel() []float32 {
}
func model() *Model {
- m := NewModel("models/base/model.onnx", "0", DefaultConfig(), true, true, false) // TODO check if CUDA is available
+ opts := Options{
+ WithCache: true,
+ WithLogits: true,
+ WithLogProbs: false,
+ }
+
+ m := NewModel("models/base/model.onnx", "0", DefaultConfig(), opts) // TODO check if CUDA is available
if err := m.Init(); err != nil {
log.Fatal(err)