diff --git a/.gitignore b/.gitignore
index acd07e1..4c2a8b7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -132,5 +132,7 @@ lightning/lightning_logs/
 results.csv
 predictions.csv
 wandb/
-lightning_logs/
+
 .cachedir/**
+lightning_logs/
+
diff --git a/run_pipeline.py b/run_pipeline.py
index 46d249e..7979a8e 100644
--- a/run_pipeline.py
+++ b/run_pipeline.py
@@ -12,6 +12,13 @@ from feature_selection.dim_reduction import reduce_dimensionality
 import ray
 ray.init()
 
+
+def run_pipeline(project_name:str, with_wandb: bool, sweep: bool):
+    wandb, model_config, training_config, data_config = setup_pipeline(project_name, with_wandb, sweep)
+    results, all_predictions, all_probabilities = run_training(project_name, wandb, sweep, model_config, training_config, data_config)  
+    reporting(results, all_predictions, all_probabilities, model_config, wandb, sweep, project_name)
+    
+
 def setup_pipeline(project_name:str, with_wandb: bool, sweep: bool):
     model_config, training_config, data_config = get_default_level_2_daily_config()
     wandb = None
@@ -20,9 +27,15 @@ def setup_pipeline(project_name:str, with_wandb: bool, sweep: bool):
         register_config_with_wandb(wandb, model_config, training_config, data_config)
     model_config, training_config, data_config = preprocess_config(model_config, training_config, data_config)
 
-    pipeline(project_name, wandb, sweep, model_config, training_config, data_config)  
 
-def pipeline(project_name:str, wandb, sweep:bool, model_config:dict, training_config:dict, data_config:dict):
+    return wandb, model_config, training_config, data_config
+    
+
+
+
+
+
+def run_training(project_name:str, wandb, sweep:bool, model_config:dict, training_config:dict, data_config:dict):
     results = pd.DataFrame()
     all_predictions = pd.DataFrame()
     all_probabilities = pd.DataFrame()
@@ -102,7 +115,10 @@ def pipeline(project_name:str, wandb, sweep:bool, model_config:dict, training_co
             all_predictions = pd.concat([all_predictions, ensemble_preds], axis=1)
             all_probabilities = pd.concat([all_probabilities, ensemble_probabilities], axis=1).fillna(0.)
         
-
+    return results, all_predictions, all_probabilities
+    
+    
+def reporting(results:pd.DataFrame, all_predictions:pd.DataFrame, all_probabilities:pd.DataFrame, model_config:dict, wandb, sweep: bool, project_name:str):
     results.to_csv('results.csv')
 
     level1_columns = results[[column for column in results.columns if 'lvl1' in column]]
@@ -134,4 +150,4 @@ def pipeline(project_name:str, wandb, sweep:bool, model_config:dict, training_co
             wandb.finish()
     
 if __name__ == '__main__':
-    setup_pipeline(project_name='price-prediction', with_wandb = False, sweep = False)
\ No newline at end of file
+    run_pipeline(project_name='price-prediction', with_wandb = False, sweep = False)
\ No newline at end of file
diff --git a/tests/test_evaluation.py b/tests/test_evaluation.py
index 529ccb8..9604234 100644
--- a/tests/test_evaluation.py
+++ b/tests/test_evaluation.py
@@ -37,6 +37,7 @@ class EvenOddStubModel(Model):
 
     data_scaling = "unscaled"
     only_column = None
+    predict_window_size = 'single_timestamp'
 
     def __init__(self, window_length) -> None:
         super().__init__()
diff --git a/tests/test_walk_forward.py b/tests/test_walk_forward.py
index cc26988..b0cb7b4 100644
--- a/tests/test_walk_forward.py
+++ b/tests/test_walk_forward.py
@@ -35,6 +35,8 @@ class IncrementingStubModel(Model):
 
     data_scaling = "unscaled"
     only_column = None
+    predict_window_size = 'single_timestamp'
+
 
     def __init__(self, window_length) -> None:
         super().__init__()
diff --git a/training/walk_forward.py b/training/walk_forward.py
index eb40c5f..b1f0f54 100644
--- a/training/walk_forward.py
+++ b/training/walk_forward.py
@@ -38,12 +38,12 @@ def walk_forward_train_test(
         scaler = clone(scaler)
 
     for index in tqdm(range(train_from, train_till)):
+        if expanding_window:
+            train_window_start = first_nonzero_return
+        else:
+            train_window_start = index - window_size - 1
 
         if iterations_before_retrain <= 0 or pd.isna(models[index-1]):
-            if expanding_window:
-                train_window_start = first_nonzero_return
-            else:
-                train_window_start = index - window_size - 1
 
             train_window_end = index - 1
             
@@ -73,7 +73,11 @@ def walk_forward_train_test(
 
         models[index] = current_model
 
-        next_timestep = X.iloc[index].to_numpy().reshape(1, -1)
+        if model.predict_window_size == 'window_size': 
+            next_timestep = X.iloc[train_window_start:index].to_numpy()#.reshape(1, -1)
+        else: 
+            next_timestep = X.iloc[index].to_numpy().reshape(1, -1)
+        
         if is_scaling_on:
             next_timestep = scaler.transform(next_timestep)