diff --git a/tests/pytest/test_classification/test_classification.py b/tests/pytest/test_classification/test_classification.py
index bcb151b1ddd042e5041ee3ff3dfa72b5a8e147c2..269603494f597107bdf589c0529d39729b288172 100644
--- a/tests/pytest/test_classification/test_classification.py
+++ b/tests/pytest/test_classification/test_classification.py
@@ -25,15 +25,18 @@ import matplotlib.pyplot as plt
 
 
 def test_sisso_classifier():
-    initialize_values_arr(80, 20, 10, 2)
+    task_sizes_train = [80]
+    task_sizes_test = [20]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 10, 2)
 
     inputs = Inputs()
     inputs.task_keys = ["task"]
     inputs.sample_ids_train = [str(ii) for ii in range(20, 100)]
     inputs.sample_ids_test = [str(ii) for ii in range(20)]
 
-    train_data = np.random.random((10, 80)) * 2.0 - 1.0
-    test_data = np.random.random((10, 20)) * 2.0 - 1.0
+    train_data = np.random.random((10, task_sizes_train[0])) * 2.0 - 1.0
+    test_data = np.random.random((10, task_sizes_test[0])) * 2.0 - 1.0
 
     train_data[0][:20] = np.random.random(20) * -1.0 - 1.0
     train_data[0][20:40] = np.random.random(20) + 1.0
@@ -105,9 +108,9 @@ def test_sisso_classifier():
     inputs.n_residual = 1
     inputs.n_models_store = 1
     inputs.task_names = ["all"]
-    inputs.task_sizes_train = [80]
-    inputs.task_sizes_test = [20]
-    inputs.leave_out_inds = list(range(20))
+    inputs.task_sizes_train = task_sizes_train
+    inputs.task_sizes_test = task_sizes_test
+    inputs.leave_out_inds = list(range(task_sizes_test[0]))
     inputs.prop_label = "Class"
 
     feat_space = FeatureSpace(inputs)
diff --git a/tests/pytest/test_descriptor_identifier/test_log_regressor.py b/tests/pytest/test_descriptor_identifier/test_log_regressor.py
index 91df5d853bc3ebc45f61f94a48eda7a27457357a..e28f250d9846b6679bdd9c44284753f026b0730c 100644
--- a/tests/pytest/test_descriptor_identifier/test_log_regressor.py
+++ b/tests/pytest/test_descriptor_identifier/test_log_regressor.py
@@ -24,22 +24,29 @@ from sissopp import (
 
 
 def test_sisso_log_regressor():
-    initialize_values_arr(90, 10, 10, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train[0], task_sizes_test[0], 10, 2)
+
     inputs = Inputs()
     inputs.phi_0 = [
         FeatureNode(
             ff,
             f"feat_{ff}",
-            np.random.random(90) * 1e2,
-            np.random.random(10) * 1e2,
+            np.random.random(task_sizes_train[0]) * 1e2,
+            np.random.random(task_sizes_test[0]) * 1e2,
             Unit(),
         )
         for ff in range(10)
     ]
 
     inputs.task_names = ["task"]
-    inputs.sample_ids_train = [str(ii) for ii in range(10, 100)]
-    inputs.sample_ids_test = [str(ii) for ii in range(10)]
+    inputs.sample_ids_train = [
+        str(ii)
+        for ii in range(task_sizes_test[0], task_sizes_test[0] + task_sizes_train[0])
+    ]
+    inputs.sample_ids_test = [str(ii) for ii in range(task_sizes_test[0])]
 
     a0 = 0.95
     a1 = 1.01
@@ -60,9 +67,9 @@ def test_sisso_log_regressor():
     inputs.n_dim = 2
     inputs.n_residual = 1
     inputs.n_models_store = 1
-    inputs.task_sizes_train = [90]
-    inputs.task_sizes_test = [10]
-    inputs.leave_out_inds = list(range(10))
+    inputs.task_sizes_train = task_sizes_train
+    inputs.task_sizes_test = task_sizes_test
+    inputs.leave_out_inds = list(range(task_sizes_test[0]))
     inputs.fix_intercept = False
     inputs.prop_label = "prop"
     inputs.prop_unit = Unit("m")
diff --git a/tests/pytest/test_descriptor_identifier/test_regressor.py b/tests/pytest/test_descriptor_identifier/test_regressor.py
index 672d94e01924a0d75a517295eef17735f7f71f13..062f0a21bce702b2283028ad40ded73b214b73cf 100644
--- a/tests/pytest/test_descriptor_identifier/test_regressor.py
+++ b/tests/pytest/test_descriptor_identifier/test_regressor.py
@@ -24,22 +24,27 @@ from sissopp import (
 
 
 def test_sisso_regressor():
-    initialize_values_arr(95, 5, 10, 2)
+    task_sizes_train = [95]
+    task_sizes_test = [5]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 10, 2)
     inputs = Inputs()
     inputs.phi_0 = [
         FeatureNode(
             ff,
             f"feat_{ff}",
-            np.random.random(95) * 1e2 - 50,
-            np.random.random(5) * 1e2 - 50,
+            np.random.random(task_sizes_train[0]) * 1e2 - 50,
+            np.random.random(task_sizes_test[0]) * 1e2 - 50,
             Unit(),
         )
         for ff in range(10)
     ]
 
     inputs.task_names = ["task"]
-    inputs.sample_ids_train = [str(ii) for ii in range(5, 100)]
-    inputs.sample_ids_test = [str(ii) for ii in range(5)]
+    inputs.sample_ids_train = [
+        str(ii)
+        for ii in range(task_sizes_test[0], task_sizes_test[0] + task_sizes_train[0])
+    ]
+    inputs.sample_ids_test = [str(ii) for ii in range(task_sizes_test[0])]
 
     a0 = np.random.random() * 5.0 - 2.5
     a1 = np.random.random() * 5.0 - 2.5
@@ -62,9 +67,9 @@ def test_sisso_regressor():
     inputs.n_dim = 2
     inputs.n_residual = 1
     inputs.n_models_store = 1
-    inputs.task_sizes_train = [95]
-    inputs.task_sizes_test = [5]
-    inputs.leave_out_inds = list(range(5))
+    inputs.task_sizes_train = task_sizes_train
+    inputs.task_sizes_test = task_sizes_test
+    inputs.leave_out_inds = list(range(task_sizes_test[0]))
     inputs.fix_intercept = False
     inputs.prop_label = "prop"
     inputs.prop_unit = Unit("m")
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_abs_diff_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_abs_diff_node.py
index 50694fcac1b2b135d63fdaa9355037ff1957d169..0fd9642cc63a8401dc37895c0d3801916774620e 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_abs_diff_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_abs_diff_node.py
@@ -28,16 +28,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_abs_diff_node():
-    initialize_values_arr(900, 10, 4, 2)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
 
-    data_1 = np.random.random(900) * 1e10 + 1e-10
-    test_data_1 = np.random.random(10) * 1e10 + 1e-10
+    initialize_values_arr(task_sizes_train, task_sizes_test, 4, 2)
 
-    data_2 = np.random.random(900) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e10 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e10 + 1e-10
 
-    data_3 = np.random.random(900) * 1e4 + 1e-10
-    test_data_3 = np.random.random(10) * 1e4 + 1e-10
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
+
+    data_3 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit("s"))
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit("m"))
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_abs_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_abs_node.py
index 58f99375a776bf89e3c752344b74d1b696488bc8..925938a4a8b26746235bab3b1cbe34e182d29ebc 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_abs_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_abs_node.py
@@ -21,15 +21,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_abs_node():
-    initialize_values_arr(90, 10, 3, 2)
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    task_sizes_train = [90]
+    task_sizes_test = [10]
 
-    data_2 = np.random.choice([1.0, -1.0], 90)
-    test_data_2 = np.random.choice([1.0, -1.0], 10)
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_3 = np.random.random(90) * 2e4 - 1e4
-    test_data_3 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.choice([1.0, -1.0], task_sizes_train[0])
+    test_data_2 = np.random.choice([1.0, -1.0], task_sizes_test[0])
+
+    data_3 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_3 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_add_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_add_node.py
index 426ca94926230b1e461e34aa130fcd2802b4b2ac..de115c458f77f3c912d6e363572eb4189cc38a99 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_add_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_add_node.py
@@ -21,16 +21,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_add_node():
-    initialize_values_arr(90, 10, 4, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    initialize_values_arr(task_sizes_train, task_sizes_test, 4, 2)
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_3 = np.random.random(90) * 1e10 + 1e-10
-    test_data_3 = np.random.random(10) * 1e10 + 1e-10
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
+
+    data_3 = np.random.random(task_sizes_train[0]) * 1e10 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 1e10 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit("s"))
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit("m"))
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_cb_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_cb_node.py
index 27e8f6a641fc645ed3321fca140fa605b79c0034..81863cf910069cf43b281f303bb34b348a589809 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_cb_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_cb_node.py
@@ -30,12 +30,15 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cube_node():
-    initialize_values_arr(90, 10, 2, 2)
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    task_sizes_train = [90]
+    task_sizes_test = [10]
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_cbrt_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_cbrt_node.py
index 1a3551ca0f73219176cf8bec8627ccd359e4b33b..0c0b2e2ff39eb5d18c7d02334b20481fddf13f86 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_cbrt_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_cbrt_node.py
@@ -31,13 +31,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cbrt_node():
-    initialize_values_arr(90, 10, 2, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_cos_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_cos_node.py
index c0d3faa2fec6314232159dda001ac9ba444f311f..da510ead267f0ed2c6681c49c188609ca10d8209 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_cos_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_cos_node.py
@@ -21,13 +21,15 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cos_node():
-    initialize_values_arr(90, 10, 3, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
 
-    data_1 = np.random.randint(0, 10000, 90) * 2.0 * np.pi
-    test_data_1 = np.random.randint(0, 10000, 10) * 2.0 * np.pi
+    data_1 = np.random.randint(0, 10000, task_sizes_train[0]) * 2.0 * np.pi
+    test_data_1 = np.random.randint(0, 10000, task_sizes_test[0]) * 2.0 * np.pi
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_div_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_div_node.py
index 216a7a9cfa1910d345a31b7fb0511f182b133c8f..07f10a89cccc713fc1c958c31122a186e98142a8 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_div_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_div_node.py
@@ -28,16 +28,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_div_node():
-    initialize_values_arr(90, 10, 5, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 5, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
-    data_3 = np.random.random(90) * 1e10 + 1e-10
-    test_data_3 = np.random.random(10) * 1e10 + 1e-10
+    data_3 = np.random.random(task_sizes_train[0]) * 1e10 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 1e10 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit("s"))
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit("m"))
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_exp_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_exp_node.py
index b1e9c4cb75bc8407bcd39ad344e65006491dfec9..670687bffdc5c7d07b6dcceb2a3c565a6459beb6 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_exp_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_exp_node.py
@@ -30,16 +30,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_exp_node():
-    initialize_values_arr(90, 10, 3, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
 
-    data_1 = np.random.random(90) + 1e-10
-    test_data_1 = np.random.random(10) + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) + 1e-10
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
-    data_3 = np.random.random(90) * 10.0 + 1e-10
-    test_data_3 = np.random.random(10) * 10.0 + 1e-10
+    data_3 = np.random.random(task_sizes_train[0]) * 10.0 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 10.0 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_inv_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_inv_node.py
index 219dbfa8d8a260aca46069397b53435f6ecefe84..01ae9dc5d3044c6d09ed5b2b3373b51b5b6fc39a 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_inv_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_inv_node.py
@@ -29,16 +29,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_inv_node():
-    initialize_values_arr(90, 10, 4, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 4, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
-    data_3 = np.random.random(90) * 1e10 + 1e-10
-    test_data_3 = np.random.random(10) * 1e10 + 1e-10
+    data_3 = np.random.random(task_sizes_train[0]) * 1e10 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 1e10 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_log_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_log_node.py
index f4643ff0356cb4fe11a383607bdba8e0722d7314..f995a41116de109832c6eb03092142de47d74dac 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_log_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_log_node.py
@@ -36,16 +36,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_log_node():
-    initialize_values_arr(90, 10, 3, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
-    data_3 = np.random.random(90) * 1.0 + 1e-10
-    test_data_3 = np.random.random(10) * 1.0 + 1e-10
+    data_3 = np.random.random(task_sizes_train[0]) * 1.0 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 1.0 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_mult_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_mult_node.py
index ed02b56753461969764b6b5c11e3869f62040ff2..8c7a9b2937cdf1ac4312b1d0fc6da32e33e5af81 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_mult_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_mult_node.py
@@ -28,16 +28,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_mult_node():
-    initialize_values_arr(90, 10, 4, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 4, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
-    data_3 = np.abs(np.random.random(90) * 1e10) + 0.1
-    test_data_3 = np.abs(np.random.random(10) * 1e10) + 0.1
+    data_3 = np.abs(np.random.random(task_sizes_train[0]) * 1e10) + 0.1
+    test_data_3 = np.abs(np.random.random(task_sizes_test[0]) * 1e10) + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit("s"))
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit("m"))
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_neg_exp_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_neg_exp_node.py
index f4d05d34408cba84c17a2730db59c50f590f28a4..2af614882498399f85612db44a9b0d05e8d8b2fd 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_neg_exp_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_neg_exp_node.py
@@ -30,16 +30,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_neg_exp_node():
-    initialize_values_arr(90, 10, 3, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
 
-    data_1 = np.random.random(90) + 1e-10
-    test_data_1 = np.random.random(10) + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) + 1e-10
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
-    data_3 = np.random.random(90) * 10.0 + 1e-10
-    test_data_3 = np.random.random(10) * 10.0 + 1e-10
+    data_3 = np.random.random(task_sizes_train[0]) * 10.0 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 10.0 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_sin_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_sin_node.py
index 1eddab2bb8b4e44c082992218018589e18478a14..9cef2b48ffcbae0f63e8812deba9b45c35c7ff4d 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_sin_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_sin_node.py
@@ -21,13 +21,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sin_node():
-    initialize_values_arr(90, 10, 3, 2)
-
-    data_1 = np.random.randint(0, 10000, 90) * (2.0) * np.pi + np.pi / 2.0
-    test_data_1 = np.random.randint(0, 10000, 10) * (2.0) * np.pi + np.pi / 2.0
-
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
+
+    data_1 = (
+        np.random.randint(0, 10000, task_sizes_train[0]) * (2.0) * np.pi + np.pi / 2.0
+    )
+    test_data_1 = (
+        np.random.randint(0, 10000, task_sizes_test[0]) * (2.0) * np.pi + np.pi / 2.0
+    )
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_six_pow_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_six_pow_node.py
index 0149f26ceadf16449de6bdcf732505e3abe42579..203022a2f73df921a9ba83ece0663c0a402aac1d 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_six_pow_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_six_pow_node.py
@@ -31,13 +31,15 @@ class InvalidFeatureMade(Exception):
 
 
 def test_six_pow_node():
-    initialize_values_arr(90, 10, 2, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
 
-    data_1 = np.random.random(90) * 1e1 + 1e-10
-    test_data_1 = np.random.random(10) * 1e1 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e1 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e1 + 1e-10
 
-    data_2 = np.random.choice([1.0, -1.0], 90)
-    test_data_2 = np.random.choice([1.0, -1.0], 10)
+    data_2 = np.random.choice([1.0, -1.0], task_sizes_train[0])
+    test_data_2 = np.random.choice([1.0, -1.0], task_sizes_test[0])
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_sq_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_sq_node.py
index 1b79f2eb6e0269200e85cdac78b8b12011a3d6ac..3c1fcf245505e46fd82e7d62dc2cd09ad484228f 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_sq_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_sq_node.py
@@ -28,13 +28,15 @@ class InvalidFeatureMade(Exception):
 
 
 def test_square_node():
-    initialize_values_arr(90, 10, 2, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.choice([1.0, -1.0], 90)
-    test_data_2 = np.random.choice([1.0, -1.0], 10)
+    data_2 = np.random.choice([1.0, -1.0], task_sizes_train[0])
+    test_data_2 = np.random.choice([1.0, -1.0], task_sizes_test[0])
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_sqrt_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_sqrt_node.py
index fdd558a836b65d46adc5237e8411cfc86eab53a4..0b225cd01d80f752f08a7ec5c4c66cf5792ae5a0 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_sqrt_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_sqrt_node.py
@@ -31,13 +31,15 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sqrt_node():
-    initialize_values_arr(90, 10, 2, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.random(90) * 2 - 1
-    test_data_2 = np.random.random(10) * 2 - 1
+    data_2 = np.random.random(task_sizes_train[0]) * 2 - 1
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2 - 1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_feature_creation/test_feat_generation/test_sub_node.py b/tests/pytest/test_feature_creation/test_feat_generation/test_sub_node.py
index 3fd14fdfd319d7275942cae6368f1ea41d17becb..46ed37ff561f85d52f950570d74b65b6b0f62454 100644
--- a/tests/pytest/test_feature_creation/test_feat_generation/test_sub_node.py
+++ b/tests/pytest/test_feature_creation/test_feat_generation/test_sub_node.py
@@ -21,16 +21,18 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sub_node():
-    initialize_values_arr(90, 10, 4, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 4, 2)
 
-    data_1 = np.random.random(90) * 1e4 + 1e-10
-    test_data_1 = np.random.random(10) * 1e4 + 1e-10
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 1e-10
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 1e-10
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
-    data_3 = np.random.random(90) * 1e10 + 1e-10
-    test_data_3 = np.random.random(10) * 1e10 + 1e-10
+    data_3 = np.random.random(task_sizes_train[0]) * 1e10 + 1e-10
+    test_data_3 = np.random.random(task_sizes_test[0]) * 1e10 + 1e-10
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit("s"))
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit("m"))
diff --git a/tests/pytest/test_feature_creation/test_feature_space/test_feature_space.py b/tests/pytest/test_feature_creation/test_feature_space/test_feature_space.py
index 754a29380d089466d9b9e8aa2bbfabe1fe04d591..a056c0bc1df3aef30dd74490f5b6101063009203 100644
--- a/tests/pytest/test_feature_creation/test_feature_space/test_feature_space.py
+++ b/tests/pytest/test_feature_creation/test_feature_space/test_feature_space.py
@@ -24,15 +24,18 @@ from sissopp import (
 
 
 def test_feature_space():
-    initialize_values_arr(90, 10, 10, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 10, 2)
 
     inputs = Inputs()
     inputs.phi_0 = [
         FeatureNode(
             ff,
             f"feat_{ff}",
-            np.random.random(90) * 1e2 - 50,
-            np.random.random(10) * 1e2 - 50,
+            np.random.random(task_sizes_train[0]) * 1e2 - 50,
+            np.random.random(task_sizes_test[0]) * 1e2 - 50,
             Unit(),
         )
         for ff in range(10)
@@ -41,7 +44,8 @@ def test_feature_space():
     inputs.prop_train = np.power(inputs.phi_0[0].value + inputs.phi_0[1].value, 2.0)
 
     inputs.allowed_ops = ["add", "sub", "mult", "sq", "cb", "sqrt", "cbrt"]
-    inputs.task_sizes_train = [90]
+    inputs.task_sizes_train = task_sizes_train
+    inputs.task_sizes_test = task_sizes_test
     inputs.calc_type = "regression"
     inputs.max_rung = 2
     inputs.n_sis_select = 10
diff --git a/tests/pytest/test_feature_creation/test_feature_space/test_gen_feature_space_from_file.py b/tests/pytest/test_feature_creation/test_feature_space/test_gen_feature_space_from_file.py
index 865cc68cd6876c4665394c388d5726c819c599f2..6bf76aa8b963659132ec181bd0528352107fa0ce 100644
--- a/tests/pytest/test_feature_creation/test_feature_space/test_gen_feature_space_from_file.py
+++ b/tests/pytest/test_feature_creation/test_feature_space/test_gen_feature_space_from_file.py
@@ -27,13 +27,16 @@ parent = pathlib.Path(__file__).parent.absolute()
 
 
 def test_gen_feature_space_from_file():
-    initialize_values_arr(90, 10, 10, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 10, 1)
     phi_0 = [
         FeatureNode(
             ff,
             f"feat_{ff}",
-            np.random.random(90) * 1e2 - 50,
-            np.random.random(10) * 1e2 - 50,
+            np.random.random(task_sizes_train[0]) * 1e2 - 50,
+            np.random.random(task_sizes_test[0]) * 1e2 - 50,
             Unit(),
         )
         for ff in range(10)
@@ -42,7 +45,7 @@ def test_gen_feature_space_from_file():
     prop = np.power(phi_0[0].value + phi_0[1].value, 2.0)
 
     feat_space = FeatureSpace(
-        f"{parent}/phi.txt", phi_0, prop, [90], "regression", 1, 1.0
+        f"{parent}/phi.txt", phi_0, prop, task_sizes_train, "regression", 1, 1.0
     )
     feat_space.sis(prop)
     assert feat_space.phi_selected[0].postfix_expr == "1|0|add|sq"
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_lorentizan.py b/tests/pytest/test_feature_creation/test_parameterize/test_lorentizan.py
index d165e0b452dead7c2a53c44cb818f79dae50f97e..14eaa0b06266e21483227c1ac5b0605975c76688 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_lorentizan.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_lorentizan.py
@@ -35,15 +35,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_lorentzian():
-    initialize_values_arr(900, 10, 1, 2)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 2)
     initialize_param_storage()
 
-    data_1 = np.linspace(-20.023658, 20.23658, 900)
-    test_data_1 = np.linspace(-19.98549, 19.08, 10)
+    data_1 = np.linspace(-20.023658, 20.23658, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.98549, 19.08, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = 21.4 / ((data_1 - 0.25) ** 2.0 + 7.1) - 1.478
-    optimizer = get_reg_optimizer([900], prop, 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 2, 2, 0.5, False)
 
     feat_node = InvParamNode(SqNode(feat_1, 2, 1e-50, 1e50), 3, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_abs.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_abs.py
index 4a265d0b84f018f8514874048df9ed73c0d51607..c7ee38f203cac3d797e28f7b3fc0568b09428951 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_abs.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_abs.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_abs_node():
-    initialize_values_arr(900, 10, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-20, 20, 900)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.linspace(-20, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = -2.3 * np.abs(1.55 * data_1 + 0.8751) - 1.2
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = AbsParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_abs_diff.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_abs_diff.py
index 2745cff34a35d76f9a1fbc5dbef7ae2867f80a48..0e27d3541acdb53bb5ce0bab12b5df525167d968 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_abs_diff.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_abs_diff.py
@@ -34,18 +34,21 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_abs_diff_node():
-    initialize_values_arr(900, 100, 2, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
 
-    data_1 = np.linspace(-20, 20, 900)
-    test_data_1 = np.linspace(-19.99, 19.99, 100)
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+
+    data_1 = np.linspace(-20, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
-    data_2 = np.linspace(-14.256, 18.6523, 900)
-    test_data_2 = np.linspace(-16.256, 17.6523, 100)
+    data_2 = np.linspace(-14.256, 18.6523, task_sizes_train[0])
+    test_data_2 = np.linspace(-16.256, 17.6523, task_sizes_test[0])
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
 
     prop = -2.3 * np.abs(data_1 - (1.5 * data_2 + 0.8751)) - 1.2
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = AbsDiffParamNode(feat_1, feat_2, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_add.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_add.py
index bc44eb695a82187d106b13a2a5373dd0d77f88e3..6d7d35ced976f031e53d077ba3d59e22f025824e 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_add.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_add.py
@@ -33,19 +33,22 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_add_node():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-20, 20, 90)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.linspace(-20, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
-    data_2 = np.linspace(-14.256, 18.6523, 90)
-    test_data_2 = np.linspace(-16.256, 17.6523, 10)
+    data_2 = np.linspace(-14.256, 18.6523, task_sizes_train[0])
+    test_data_2 = np.linspace(-16.256, 17.6523, task_sizes_test[0])
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
 
     prop = -2.3 * (data_1 + 1.5 * data_2) - 1.2
-    optimizer = get_reg_optimizer([90], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = AddParamNode(feat_1, feat_2, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_cb.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_cb.py
index 16ba2772a59fe4130e2b48edc408ddb2e606366b..7bd015a0213388f23d347e426205a96d254eaf1c 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_cb.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_cb.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_cb_node():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-15, 15, 90)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.linspace(-15, 15, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = 1.55 * np.power(data_1 + 0.8751, 3.0) - 1.2
-    optimizer = get_reg_optimizer([90], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = CbParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_cbrt.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_cbrt.py
index 1b0bbe60356f7753eb5aa151867c89312a22bde0..861920f9c4d511011a480ae05cb2d2d612576364 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_cbrt.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_cbrt.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_cbrt_node():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(0.5, 20, 90)
-    test_data_1 = np.linspace(0.52145, 19.99, 10)
+    data_1 = np.linspace(0.5, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(0.52145, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = np.cbrt(1.55 * data_1 + 0.8751) - 1.2
-    optimizer = get_reg_optimizer([90], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = CbrtParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_cos.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_cos.py
index 527e6978cfe175782c46f53cd837e266bee4a686..0e57422251122c8609fc9853eac10a0370a00d94 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_cos.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_cos.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_cos_node():
-    initialize_values_arr(900, 100, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-2.0 * np.pi, 2.0 * np.pi, 900)
-    test_data_1 = np.linspace(-1.994 * np.pi, 1.994 * np.pi, 100)
+    data_1 = np.linspace(-2.0 * np.pi, 2.0 * np.pi, task_sizes_train[0])
+    test_data_1 = np.linspace(-1.994 * np.pi, 1.994 * np.pi, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = -1.1 * np.cos(1.25 * data_1 + 2.13) + 0.01578
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = CosParamNode(feat_1, 1, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_div.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_div.py
index 2bbc8fbbd2f55ed0c4aae831f09d1ebf5b74e44c..be4dbc1c320dbef67732f2b227a74456bc246e66 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_div.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_div.py
@@ -34,19 +34,22 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_div_node():
-    initialize_values_arr(900, 10, 2, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
 
-    data_1 = np.random.uniform(-2.50, 2.50, 900)
-    test_data_1 = np.linspace(0.52145, 19.99, 10)
+    data_1 = np.random.uniform(-2.50, 2.50, task_sizes_train[0])
+    test_data_1 = np.linspace(0.52145, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
-    data_2 = np.random.uniform(0.5, 5.0, 900)
-    test_data_2 = np.linspace(0.41, 19.8432, 10)
+    data_2 = np.random.uniform(0.5, 5.0, task_sizes_train[0])
+    test_data_2 = np.linspace(0.41, 19.8432, task_sizes_test[0])
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
 
     prop = 4.124 * data_1 / ((data_2 + 1.8751)) - 0.12
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = DivParamNode(feat_1, feat_2, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_exp.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_exp.py
index 0cccc1b361fd900a592e72afecb913487fe5d1af..61c17fdab5950caa1ebd8799aa66d82aeabfd95e 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_exp.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_exp.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_exp_node():
-    initialize_values_arr(900, 10, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.random.uniform(-2.0, 2.0, 900)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.random.uniform(-2.0, 2.0, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = np.exp(1.05 * data_1 + 0.08751) - 0.12
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = ExpParamNode(feat_1, 1, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_inv.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_inv.py
index 2da9ff2afc3bcf24c42d99999ef86519ba19d3bd..492984bcc725b0927e37ddc7d17f381dcb8cf9ce 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_inv.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_inv.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_inv_node():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(0.5, 20, 90)
-    test_data_1 = np.linspace(1.0, 19.99, 10)
+    data_1 = np.linspace(0.5, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(1.0, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = 1.0 / (1.55 * data_1 + 0.8751) - 1.2
-    optimizer = get_reg_optimizer([90], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = InvParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_log.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_log.py
index 727d1534a1d76c7f500df9025471df08fa4708f0..33195bcbef5f5e492ddb477b0bdfbddac53cc95d 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_log.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_log.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_log_node():
-    initialize_values_arr(900, 10, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(0.5, 20, 900)
-    test_data_1 = np.linspace(0.52145, 19.99, 10)
+    data_1 = np.linspace(0.5, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(0.52145, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = -2.014 * np.log(1.15 * data_1 + 0.1387)
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = LogParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_neg_exp.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_neg_exp.py
index c867a5657f1f5faae25ca4e3e9ea6ef0b0060ac3..6d8fcda8af0d1cd5cd28c07c9c87e7bdac144b55 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_neg_exp.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_neg_exp.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_neg_exp_node():
-    initialize_values_arr(900, 10, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.random.uniform(-5.0, 5.0, 900)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.random.uniform(-5.0, 5.0, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = np.exp(-1.05 * data_1 + 0.08751) - 0.12
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = NegExpParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_sin.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_sin.py
index 32ce702e20cc88050e44fbab132c3e857705f06a..5de305a9db4b5d88365194643d36df5e9f8c5e9f 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_sin.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_sin.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_sin_node():
-    initialize_values_arr(900, 100, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-2.0 * np.pi, 2.0 * np.pi, 900)
-    test_data_1 = np.linspace(-1.994 * np.pi, 1.994 * np.pi, 100)
+    data_1 = np.linspace(-2.0 * np.pi, 2.0 * np.pi, task_sizes_train[0])
+    test_data_1 = np.linspace(-1.994 * np.pi, 1.994 * np.pi, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = -1.1 * np.sin(1.25 * data_1 + 2.13) + 0.01578
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = SinParamNode(feat_1, 1, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_six_pow.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_six_pow.py
index 20ed0986d0b11a76a59f5aa89a4e738a0936f5bb..2a2376dc24fd122c2a769c5a3e3ed089df0ed7cd 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_six_pow.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_six_pow.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_six_pow_node():
-    initialize_values_arr(900, 10, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.random.uniform(-2.0, 2.0, 900)
-    test_data_1 = np.random.uniform(-19.99, 19.99, 10)
+    data_1 = np.random.uniform(-2.0, 2.0, task_sizes_train[0])
+    test_data_1 = np.random.uniform(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = 1.55 * np.power(data_1 + 0.21, 6.0) - 0.12
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = SixPowParamNode(feat_1, 1, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_sq.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_sq.py
index 1e4b67c6c1edad2c3b36959e061a3017ae048801..50102b38df4c0f171ae46748ae929026bfce9dde 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_sq.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_sq.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_sq_node():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-20, 20, 90)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.linspace(-20, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = np.power(1.55 * data_1 + 0.8751, 2.0) - 1.2
-    optimizer = get_reg_optimizer([90], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = SqParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_sqrt.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_sqrt.py
index 3542d830f024c216515efddd8e4bb93395c4bf8f..f5144c7c3c81cf0332cd5593d2875b254b448c69 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_sqrt.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_sqrt.py
@@ -34,15 +34,18 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_sqrt_node():
-    initialize_values_arr(900, 10, 1, 1)
+    task_sizes_train = [900]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(0.5, 500, 900)
-    test_data_1 = np.linspace(0.52145, 19.99, 10)
+    data_1 = np.linspace(0.5, 500, task_sizes_train[0])
+    test_data_1 = np.linspace(0.52145, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
     prop = np.sqrt(1.55 * data_1 + 0.8751) - 1.2
-    optimizer = get_reg_optimizer([900], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = SqrtParamNode(feat_1, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_feature_creation/test_parameterize/test_param_sub.py b/tests/pytest/test_feature_creation/test_parameterize/test_param_sub.py
index f0fd0edc1892aaa3cfea7bcad7e544e1ced44671..615a76cf6a4f3ab0e2e79d9a38b14bf4aa154640 100644
--- a/tests/pytest/test_feature_creation/test_parameterize/test_param_sub.py
+++ b/tests/pytest/test_feature_creation/test_parameterize/test_param_sub.py
@@ -33,19 +33,22 @@ def check_feat_parmeters(feat, prop):
 
 
 def test_param_sub_node():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
 
-    data_1 = np.linspace(-20, 20, 90)
-    test_data_1 = np.linspace(-19.99, 19.99, 10)
+    data_1 = np.linspace(-20, 20, task_sizes_train[0])
+    test_data_1 = np.linspace(-19.99, 19.99, task_sizes_test[0])
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
 
-    data_2 = np.linspace(-14.256, 18.6523, 90)
+    data_2 = np.linspace(-14.256, 18.6523, task_sizes_train[0])
     test_data_2 = np.linspace(-16.256, 17.6523, 10)
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
 
     prop = -2.3 * (data_1 - 1.5 * data_2) - 1.2
-    optimizer = get_reg_optimizer([90], prop, 1, 1, 0.5, False)
+    optimizer = get_reg_optimizer([task_sizes_train[0]], prop, 1, 1, 0.5, False)
 
     feat_node = SubParamNode(feat_1, feat_2, 2, 1e-50, 1e50)
     feat_node.get_parameters(optimizer)
diff --git a/tests/pytest/test_model_eval/test_model_node/test_abs.py b/tests/pytest/test_model_eval/test_model_node/test_abs.py
index 95cc923023b064afd133db472d2c33ad39936f4d..ecb579b647af6c76392ba162e57ceb6ca40ac1de 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_abs.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_abs.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_abs_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_abs_diff.py b/tests/pytest/test_model_eval/test_model_node/test_abs_diff.py
index b24c4b1561e4ab38e5fab4baecf28d9db508e944..a1ea85b88b1e651ebfb4c36e15416d4d8a8ab4ef 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_abs_diff.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_abs_diff.py
@@ -27,12 +27,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_abs_diff_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
-
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_model_node/test_add.py b/tests/pytest/test_model_eval/test_model_node/test_add.py
index 44d952859a4344929551274e9173344fc2e5de25..d38e915d67e4ac79248414778033c14ea19384fa 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_add.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_add.py
@@ -27,12 +27,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_add_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
-
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_model_node/test_binary_binary.py b/tests/pytest/test_model_eval/test_model_node/test_binary_binary.py
index df10363c3b299ad4bcefe1371e6807fa63fbc3e5..899b218c333b7f00cef67fc7e43abf9f5a0c879d 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_binary_binary.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_binary_binary.py
@@ -28,15 +28,17 @@ class InvalidFeatureMade(Exception):
 
 
 def test_bin_bin_model_eval():
-    initialize_values_arr(90, 10, 3, 2)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
-    data_3 = np.random.random(90) * 2e4 - 1e4
-    test_data_3 = np.random.random(10) * 2e4 - 1e4
+    data_3 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_3 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
@@ -46,9 +48,9 @@ def test_bin_bin_model_eval():
     node_2 = AddNode(node_1, feat_2, 4, 1e-50, 1e50)
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    data_3 = np.random.random(90) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    data_3 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
     val_data = data_1 * data_3 + data_2
 
     data_dict = {"t_a": data_1[0], "x_a": data_2[0], "m_a": data_3[0]}
diff --git a/tests/pytest/test_model_eval/test_model_node/test_binary_unary.py b/tests/pytest/test_model_eval/test_model_node/test_binary_unary.py
index 0d66d222f5b65215d1d301e521976c80a8a26c46..f3fa36214226931f4f32fd4565f6f90b8d83fa9c 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_binary_unary.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_binary_unary.py
@@ -28,12 +28,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_bin_un_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
@@ -42,8 +44,8 @@ def test_bin_un_model_eval():
     node_2 = AddNode(node_1, feat_2, 2, 1e-50, 1e50)
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    data_2 = np.random.random(90) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
     val_data = data_1 ** 2.0 + data_2
 
     data_dict = {"t_a": data_1[0], "x_a": data_2[0]}
diff --git a/tests/pytest/test_model_eval/test_model_node/test_cb.py b/tests/pytest/test_model_eval/test_model_node/test_cb.py
index d1fc6903f9dab100cc511c171ec95212ff93a8ac..d989a5ceb2721a509fa5310d6560aa89879aa1af 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_cb.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_cb.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cb_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_cbrt.py b/tests/pytest/test_model_eval/test_model_node/test_cbrt.py
index e38e4465b87460bebd8717ff808e95af29bd8eea..c833ae04672e1291fb7b93d3dfec14f38b5f3ae2 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_cbrt.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_cbrt.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cbrt_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_cos.py b/tests/pytest/test_model_eval/test_model_node/test_cos.py
index 42e2a44ea6c474a9b23a70ba595e563d0139d5bd..d75e39119d7c02d7173fe2dc323f8804dd398efd 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_cos.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_cos.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cos_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_div.py b/tests/pytest/test_model_eval/test_model_node/test_div.py
index 2b9549e543968c7785cebfabc78cf69fef2b65aa..31dac1a4e0308f4755db90afe4c3da6b3ed9229f 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_div.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_div.py
@@ -27,12 +27,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_div_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
-
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_model_node/test_exp.py b/tests/pytest/test_model_eval/test_model_node/test_exp.py
index 76c0cc9c9c90f4c382f0445f043a642c1e99a7b1..93bf802d36ecb464631f54ef42f23fd6a799c9b8 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_exp.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_exp.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_exp_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e1 - 1e1
-    test_data_1 = np.random.random(10) * 2e1 - 1e1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e1 - 1e1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e1 - 1e1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_inv.py b/tests/pytest/test_model_eval/test_model_node/test_inv.py
index 543d4f0b2f4b473b8f7c5434d7675069362ca132..34e193e3863544e1f09d8e0b10fecec946635ef4 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_inv.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_inv.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_inv_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_log.py b/tests/pytest/test_model_eval/test_model_node/test_log.py
index bcf597b60005f994368eea06a501903ace40c6c9..a29dbcf6f317b70e22e7616ceebf181f66c8b636 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_log.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_log.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_log_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_mult.py b/tests/pytest/test_model_eval/test_model_node/test_mult.py
index 53d995c009a0f99088904143497c06567847cf50..7c70dc8556059fbd3f689a61f46955d574b5b717 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_mult.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_mult.py
@@ -27,12 +27,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_mult_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
-
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_model_node/test_neg_exp.py b/tests/pytest/test_model_eval/test_model_node/test_neg_exp.py
index b8d6deee8f03835a53cfef099565b837840a0cb7..bb0e844b9b40c0e5b1ba02ca5e24dd4051734c62 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_neg_exp.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_neg_exp.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_neg_exp_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e1 - 1e1
-    test_data_1 = np.random.random(10) * 2e1 - 1e1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e1 - 1e1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e1 - 1e1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_sin.py b/tests/pytest/test_model_eval/test_model_node/test_sin.py
index f52212fecca55edd5d52472b5ccbffc60f8228dc..a08b3fe49104461ca20a1bb0e403c2f6d22a80a5 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_sin.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_sin.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sin_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_six_pow.py b/tests/pytest/test_model_eval/test_model_node/test_six_pow.py
index 1e7f5fd9e05523e20f8a90f079d8163a9a8b59dd..19d20f433a902a0f1fa56e1d09e2aec6ec1d645b 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_six_pow.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_six_pow.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_six_pow_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_sq.py b/tests/pytest/test_model_eval/test_model_node/test_sq.py
index 0ec268f658845187d0a090347ca24a0bb4a91030..a47f9ebf98b3cd5c5738509686cab13fa1f46e41 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_sq.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_sq.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sq_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_sqrt.py b/tests/pytest/test_model_eval/test_model_node/test_sqrt.py
index 3c839a90623500e7e13dd396cb6d0fae10575d2f..4d10d6bb7d4cbedc26d0a3822bee66597da2599c 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_sqrt.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_sqrt.py
@@ -27,9 +27,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sqrt_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_model_node/test_sub.py b/tests/pytest/test_model_eval/test_model_node/test_sub.py
index 95398958efdee70559d0af65100b7af50fb53f8b..92cecc3c20773d61bf6044301b647aaa6f9e16f1 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_sub.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_sub.py
@@ -27,12 +27,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sub_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
-
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
+
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_model_node/test_unary_binary.py b/tests/pytest/test_model_eval/test_model_node/test_unary_binary.py
index b32b610d5bb940afb491fab2fde1e8c9b93146e4..1ec0e4dc0a1cec642ece9f180f9bef68882cc7f7 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_unary_binary.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_unary_binary.py
@@ -28,12 +28,14 @@ class InvalidFeatureMade(Exception):
 
 
 def test_un_bin_model_eval():
-    initialize_values_arr(90, 10, 2, 2)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
@@ -42,8 +44,8 @@ def test_un_bin_model_eval():
     node_2 = SqNode(node_1, 3, 1e-50, 1e50)
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    data_2 = np.random.random(90) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
     val_data = (data_1 + data_2) ** 2.0
 
     data_dict = {"t_a": data_1[0], "x_a": data_2[0]}
diff --git a/tests/pytest/test_model_eval/test_model_node/test_unary_unary.py b/tests/pytest/test_model_eval/test_model_node/test_unary_unary.py
index a9e94295788a3b046772d1fffcb3f576f33f53de..79a5a994b13715a930524db3da9161da32b3a622 100644
--- a/tests/pytest/test_model_eval/test_model_node/test_unary_unary.py
+++ b/tests/pytest/test_model_eval/test_model_node/test_unary_unary.py
@@ -28,9 +28,11 @@ class InvalidFeatureMade(Exception):
 
 
 def test_un_un_model_eval():
-    initialize_values_arr(90, 10, 1, 2)
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 2)
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
@@ -39,7 +41,7 @@ def test_un_un_model_eval():
     node_2 = InvNode(node_1, 3, 1e-50, 1e50)
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
     val_data = 1.0 / data_1 ** 2.0
 
     data_dict = {"t_a": data_1[0]}
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_abs_diff_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_abs_diff_param.py
index 35a14991f7b0953bc223decab1fc27bb5d44727e..9ff404aa5f1b455b2c4db31a111555251cb594ac 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_abs_diff_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_abs_diff_param.py
@@ -29,15 +29,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_abs_diff_param_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_abs_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_abs_param.py
index 2e8069fdbb4c6c9fd8146c45aa63092430708cf6..06906b5f0381b8720d50828dbc91623d2d71780c 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_abs_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_abs_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_abs_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_add_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_add_param.py
index 7b0c321e802eac5de00721e246ecce38a9328c7f..44e8a998787da24ce91eafb58cba728108a35df5 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_add_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_add_param.py
@@ -29,15 +29,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_add_param_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_binary_binary_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_binary_binary_param.py
index 4077bccb6d59af224344979ddd82de99b6b3efa5..82c32af8f63f1acc866f0323cc68a1328c97fd25 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_binary_binary_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_binary_binary_param.py
@@ -30,18 +30,22 @@ class InvalidFeatureMade(Exception):
 
 
 def test_bin_bin_model_eval():
-    initialize_values_arr(90, 10, 3, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 3, 2)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
-    data_3 = np.random.random(90) * 2e4 - 1e4
-    test_data_3 = np.random.random(10) * 2e4 - 1e4
+    data_3 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_3 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
@@ -55,9 +59,9 @@ def test_bin_bin_model_eval():
 
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    data_3 = np.random.random(90) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    data_3 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
     val_data = (data_1 * (3.0 * data_3 + 4.0)) + data_2 + 2.0
 
     data_dict = {"t_a": data_1[0], "x_a": data_2[0], "m_a": data_3[0]}
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_binary_unary_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_binary_unary_param.py
index a95070ae8c1e60ce7163938ff489f8e84843b818..dec60fdaf88e9c2be517537de094fa94b81386d9 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_binary_unary_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_binary_unary_param.py
@@ -30,15 +30,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_bin_un_model_eval():
-    initialize_values_arr(90, 10, 2, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
@@ -51,8 +55,8 @@ def test_bin_un_model_eval():
 
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    data_2 = np.random.random(90) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
     val_data = (3.0 * data_1 + 4.0) ** 2.0 + data_2 + 2.0
 
     data_dict = {"t_a": data_1[0], "x_a": data_2[0]}
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_cb_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_cb_param.py
index 09dbd11d6707f07c70b3bec0453fc3e0dd5b6033..7849eca5fa748b17c1790571ef5b0a6b81cba71e 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_cb_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_cb_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cb_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_cbrt_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_cbrt_param.py
index b04a6087567fad307a19f0c91e86658049cab8c7..6b476c57b5de487c6371cc38130f68aea7cd718d 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_cbrt_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_cbrt_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cbrt_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_cos_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_cos_param.py
index b8f7661ed70a2a2c6096aefab69366b30724c2cf..0b7a622361869579d866166933b17cd87a427602 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_cos_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_cos_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_cos_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_div_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_div_param.py
index 8f5d19ca589d61e088967c4bc1652d59a7546864..7e119467b3ba522e486f0f9bfc5ecd4d814d8294 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_div_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_div_param.py
@@ -29,15 +29,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_div_param_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_exp_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_exp_param.py
index b6dd4de2069186997b1bc12205c8aef8283e2066..f09f45b408f0f36ce78d8b7c226620726dfea759 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_exp_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_exp_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_exp_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e1 - 1e1
-    test_data_1 = np.random.random(10) * 2e1 - 1e1
+    data_1 = np.random.random(task_sizes_train[0]) * 2e1 - 1e1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e1 - 1e1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_inv_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_inv_param.py
index ee92eb5a0e7fe72866d7814af8166b26dfc3f6cd..c745617435246e172a4ef2e64d6e6a682b4867bc 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_inv_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_inv_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_inv_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_log_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_log_param.py
index accfb658d3e5f3a319b45c6d8a92cbefe6c70c20..ea8a5f129285da24fa89d9b82f667ad84356b6d9 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_log_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_log_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_log_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_mult_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_mult_param.py
index b66c259e2c0c428779030da2957ff78fa20edd89..6395b2c4014318eff7ff3de24618859bc0264b8d 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_mult_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_mult_param.py
@@ -29,15 +29,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_mult_param_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_neg_exp_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_neg_exp_param.py
index 50dcd99c0239b60504d0dacd27c6e5483ef9911a..1ae976bab60418e38364d49d1004ab6e459b9549 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_neg_exp_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_neg_exp_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_neg_exp_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e1 - 1e1
-    test_data_1 = np.random.random(10) * 2e1 - 1e1
+    data_1 = np.random.random(task_sizes_train[0]) * 2e1 - 1e1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e1 - 1e1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_sin_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_sin_param.py
index 19ef2d0b917df9c2bdf6501355037f182772b447..028082b30c458027a40722324f34e0b607b07c20 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_sin_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_sin_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sin_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_six_pow_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_six_pow_param.py
index d282cbf21ad071657b88590db3838fd6fc7ec850..5d6b349b859bd58a169c8c1bc3d2a0f09d71888f 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_six_pow_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_six_pow_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_six_pow_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_sq_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_sq_param.py
index 36d8ad1bf4707dacb7eb28789f8ac4c79dac63fc..41c21e63a41b1605fe648089e0e98b471833133f 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_sq_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_sq_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sq_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 2e4 - 1e4
-    test_data_1 = np.random.random(10) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_1 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_sqrt_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_sqrt_param.py
index e18a1750ed8d419812820bae17e6b04a6231f3e0..bb47b82245b4ed24b54a87c408b11a2e34343d41 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_sqrt_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_sqrt_param.py
@@ -29,12 +29,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sqrt_param_model_eval():
-    initialize_values_arr(90, 10, 1, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_sub_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_sub_param.py
index abe0cdb3932ef284a7ab48d72f7a03b5d6454f2c..d7b12d0d8a421444735498e704816c56247ae6da 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_sub_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_sub_param.py
@@ -29,15 +29,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_sub_parm_model_eval():
-    initialize_values_arr(90, 10, 2, 1)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 1)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_unary_binary_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_unary_binary_param.py
index 9f828d3934e39fd18481687fb89e27ee6e01687a..2d0eb6077491a25368178fcb2c12941d0e280f23 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_unary_binary_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_unary_binary_param.py
@@ -30,15 +30,19 @@ class InvalidFeatureMade(Exception):
 
 
 def test_un_bin_model_eval():
-    initialize_values_arr(90, 10, 2, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 2, 2)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
-    data_2 = np.random.random(90) * 2e4 - 1e4
-    test_data_2 = np.random.random(10) * 2e4 - 1e4
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
+    test_data_2 = np.random.random(task_sizes_test[0]) * 2e4 - 1e4
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feat_2 = FeatureNode(1, "x_a", data_2, test_data_2, Unit())
@@ -51,8 +55,8 @@ def test_un_bin_model_eval():
 
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    data_2 = np.random.random(90) * 2e4 - 1e4
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    data_2 = np.random.random(task_sizes_train[0]) * 2e4 - 1e4
     val_data = ((data_1 + 0.5 * data_2 + 4.0) + 2.0) ** 2.0
 
     data_dict = {"t_a": data_1[0], "x_a": data_2[0]}
diff --git a/tests/pytest/test_model_eval/test_param_model_node/test_unary_unary_param.py b/tests/pytest/test_model_eval/test_param_model_node/test_unary_unary_param.py
index 85f879dfbf03c6aad732aa8384a9f78443e1a171..e2fb09237ddd55f809f072c033657d3e0f03afeb 100644
--- a/tests/pytest/test_model_eval/test_param_model_node/test_unary_unary_param.py
+++ b/tests/pytest/test_model_eval/test_param_model_node/test_unary_unary_param.py
@@ -30,12 +30,16 @@ class InvalidFeatureMade(Exception):
 
 
 def test_un_un_model_eval():
-    initialize_values_arr(90, 10, 1, 2)
+    task_sizes_train = [90]
+    task_sizes_test = [10]
+    initialize_values_arr(task_sizes_train, task_sizes_test, 1, 2)
     initialize_param_storage()
-    optimizer = get_reg_optimizer([90], np.zeros(90), 2, 2, 0.5, False)
+    optimizer = get_reg_optimizer(
+        [task_sizes_train[0]], np.zeros(task_sizes_train[0]), 2, 2, 0.5, False
+    )
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
-    test_data_1 = np.random.random(10) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
+    test_data_1 = np.random.random(task_sizes_test[0]) * 1e4 + 0.1
 
     feat_1 = FeatureNode(0, "t_a", data_1, test_data_1, Unit())
     feats = []
@@ -48,7 +52,7 @@ def test_un_un_model_eval():
 
     model_node = ModelNode(node_2)
 
-    data_1 = np.random.random(90) * 1e4 + 0.1
+    data_1 = np.random.random(task_sizes_train[0]) * 1e4 + 0.1
     val_data = 1.0 / ((3.0 * data_1 + 4.0) ** 2.0 + 2.0)
 
     data_dict = {"t_a": data_1[0]}
diff --git a/tests/pytest/test_sisso.py b/tests/pytest/test_sisso.py
index 4a0401ba6ca7d5c94d43739c647fd71c0f1e3538..dbbc9c4b358934df4377f4a7bf5be1709ab30b0e 100644
--- a/tests/pytest/test_sisso.py
+++ b/tests/pytest/test_sisso.py
@@ -37,8 +37,8 @@ def test_sisso():
 
     sisso.fit()
 
-    # shutil.rmtree("models/")
-    # shutil.rmtree("feature_space/")
+    shutil.rmtree("models/")
+    shutil.rmtree("feature_space/")
 
     assert sisso.models[1][0].rmse < 1e-7
     assert sisso.models[1][0].test_rmse < 1e-7