Save segments after actually segmenting

1 files changed, 15 insertions, 9 deletions

diff --git a/research/knobloch/cmd/train/main.go b/research/knobloch/cmd/train/main.go
index 5436784..72f7e00 100644
--- a/research/knobloch/cmd/train/main.go
+++ b/research/knobloch/cmd/train/main.go
@@ -107,26 +107,32 @@ func train() {
 
 		fmt.Printf("%s\n\n", c.string)
 
+		saveSegments := false
+
 		if fileOpen, errOpen := os.Open(filepath.Join(out, "segments.gob")); errOpen != nil {
+			saveSegments = true
+		} else {
+			defer fileOpen.Close()
+
+			if err := t.LoadSegments(fileOpen); err != nil {
+				log.Fatal(err)
+			}
+		}
+
+		t.Train()
+
+		if saveSegments {
 			if fileCreate, errCreate := os.Create(filepath.Join(out, "segments.gob")); errCreate != nil {
 				log.Fatal(errCreate)
 			} else {
 				defer fileCreate.Close()
 
 				if err := t.SaveSegments(fileCreate); err != nil {
-					log.Fatal(errOpen)
+					log.Fatal(err)
 				}
 			}
-		} else {
-			defer fileOpen.Close()
-
-			if err := t.LoadSegments(fileOpen); err != nil {
-				log.Fatal(errOpen)
-			}
 		}
 
-		t.Train()
-
 		dir := filepath.Join(out, c.string)
 
 		if err := os.Mkdir(dir, 0755); err != nil && !errors.Is(err, os.ErrExist) {