diff --git a/ANN_Data_Generator.py b/ANN_Data_Generator.py
index 37f5b44675ddd671fc625273f78ca34561c13060..1dce63c21a6331713d2cb141dd8f80768276ee4a 100644
--- a/ANN_Data_Generator.py
+++ b/ANN_Data_Generator.py
@@ -12,6 +12,7 @@ TODO: Improve function selection (more even distribution) -> Done
 TODO: Add documentation -> Done
 TODO: Improve comments -> Done
 TODO: Remove unnecessary lines -> Done
+TODO: Fix bug in 'input_data' calculation -> Done
 
 """
 
@@ -205,7 +206,8 @@ class TrainingDataGenerator(object):
                 'NoDetection', polynomial_degree=polynomial_degree,
                 num_grid_cells=self._stencil_length, left_bound=left_bound, right_bound=right_bound,
                 quadrature='Gauss', quadrature_config={'num_eval_points': polynomial_degree+1})
-            input_data[i] = dg_scheme.build_training_data(adjustment, initial_condition)
+            input_data[i] = dg_scheme.build_training_data(adjustment, self._stencil_length,
+                                                          initial_condition)
 
             count += 1
             if count % 1000 == 0: