diff --git a/ANN_Data_Generator.py b/ANN_Data_Generator.py
index 07a7d7c509d1d3257398939517cdbd2f9dfe6187..e9a9b8c2a72e3dfda2184c15bde30b9b1d36ae62 100644
--- a/ANN_Data_Generator.py
+++ b/ANN_Data_Generator.py
@@ -16,6 +16,7 @@ TODO: Fix bug in 'input_data' calculation -> Done
 TODO: Change order of imports -> Done
 TODO: Improve calculation of normalization -> Done
 TODO: Improve calculation of random point for stencil -> Done
+TODO: Improve directory selection -> Done
 
 """
 
@@ -47,7 +48,7 @@ class TrainingDataGenerator(object):
 
     """
     def __init__(self, initial_conditions, left_bound=-1, right_bound=1, balance=0.5,
-                 stencil_length=3, directory=None):
+                 stencil_length=3, directory='test_data'):
         """Initializes TrainingDataGenerator.
 
         Parameters
@@ -85,9 +86,7 @@ class TrainingDataGenerator(object):
                 self._troubled_functions.append(function)
 
         # Set directory
-        self._data_dir = 'test_data'
-        if directory is not None:
-            self._data_dir = directory
+        self._data_dir = directory
         if not os.path.exists(self._data_dir):
             os.makedirs(self._data_dir)