diff --git a/Code/python/notebooks/predict.py b/Code/python/notebooks/predict.py
index 04dfefc..41238a7 100644
--- a/Code/python/notebooks/predict.py
+++ b/Code/python/notebooks/predict.py
@@ -45,10 +45,10 @@ file_list = load_file_list(file_list_path)
 
 base_ds = dataloader.EvenlySpacedDataset(
     filepaths=file_list,
-    n_input=30 * 60,
-    n_output=30 * 60,
+    n_input=warm,
+    n_output=pred,
     n_windows_per_file=5,
-    step=1,
+    step=step,
     feature_columns=("lat", "lon", "alt", "ias"),
     context_columns=("last_lat", "last_lon", "last_alt", "last_ias", "last_timestamp"),
     time_columns=("timestamp", "dt"),