Setup for new benchmark and bugfix

2 files changed, 7 insertions, 4 deletions

diff --git a/do_ml.py b/do_ml.py
index 87e8b72db..a424a610a 100644
--- a/do_ml.py
+++ b/do_ml.py
@@ -68,7 +68,10 @@ def do_ml(desc_data,
         return None
 
     if not test_size:
-        test_size = d_len - training_size
+        if d_len - training_size > 1500:
+            test_size = 1500
+        else:
+            test_size = d_len - training_size
 
     tic = time.perf_counter()
     if show_msgs:
diff --git a/main.py b/main.py
index 059fa2213..9144099d5 100644
--- a/main.py
+++ b/main.py
@@ -67,17 +67,17 @@ def main():
     procs = []
     cm_pipes = []
     ljm_pipes = []
-    for i in range(500, 1500 + 1, 500):
+    for i in range(2500, 6000 + 1, 500):
         cm_recv, cm_send = Pipe(False)
         p1 = Process(target=do_ml,
-                     args=(cm_data, energy_pbe0, i, 'CM', cm_send, 500))
+                     args=(cm_data, energy_pbe0, i, 'CM', cm_send))
         procs.append(p1)
         cm_pipes.append(cm_recv)
         p1.start()
 
         ljm_recv, ljm_send = Pipe(False)
         p2 = Process(target=do_ml,
-                     args=(ljm_data, energy_pbe0, i, 'L-JM', ljm_send, 500))
+                     args=(ljm_data, energy_pbe0, i, 'L-JM', ljm_send))
         procs.append(p2)
         ljm_pipes.append(ljm_recv)
         p2.start()