Handle environment initialization on package level

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}