summaryrefslogtreecommitdiff
path: root/gpt2/model_test.go
diff options
context:
space:
mode:
Diffstat (limited to 'gpt2/model_test.go')
-rw-r--r--gpt2/model_test.go14
1 files changed, 14 insertions, 0 deletions
diff --git a/gpt2/model_test.go b/gpt2/model_test.go
index d197ccc..0072254 100644
--- a/gpt2/model_test.go
+++ b/gpt2/model_test.go
@@ -9,6 +9,20 @@ import (
"testing"
)
+func TestMain(m *testing.M) {
+ if err := InitializeEnvironment(); err != nil {
+ log.Fatal(err)
+ }
+
+ exit := m.Run()
+
+ if err := DestroyEnvironment(); err != nil {
+ log.Fatal(err)
+ }
+
+ os.Exit(exit)
+}
+
func fromModel() []float32 {
prompt := []int64{464, 2068, 7586, 21831, 18045, 625, 262, 16931, 3290}