diff --git a/src/main.py b/src/main.py
index e926076a5d1fbfd686169f5bf3e51cee5eea75ec..a607170e4936b0b7d462b6565abe3e73426016a5 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,5 +1,5 @@
 from pathlib import Path
-from pandas import read_csv
+from pandas import read_csv, DataFrame
 from autogluon.tabular import TabularDataset, TabularPredictor
 from datetime import datetime
 
@@ -53,9 +53,9 @@ def handlePredict(ds):
 
 def predict(ds, model):
     root = Path(__file__).parent.parent
-    testing = root / "resources" / f"Testing{datasets[ds]}.csv"
+    testing = root / "resources" / "TestingData{}.csv".format("Multi" if ds =='b' else "Binary")
     testingData = read_csv(testing, header=None)
-    savePath = root / "output" / f"Testing{datasets[ds]}.csv"
+    savePath = root / "output" / "TestingResults{}.csv".format("Multi" if ds =='b' else "Binary")
 
     predictor = TabularPredictor.load(path=model)
 
@@ -63,10 +63,12 @@ def predict(ds, model):
     testingData[128] = a
     print(a)
 
-    a = savePath.absolute().as_posix()
-    print(f"Saving to {a}")
-    testingData.to_csv(path_or_buf=a, header=False)
+    savePath = savePath.absolute().as_posix()
+    print(f"Saving to {savePath}")
 
+    with open(savePath, "w") as f:
+        for i in a:
+            f.write(f"{i}\n")
 
 def handleTrain(ds):
     print("Starting training")