Update DataGenerator.py

DataGenerator.py

@@ -18,8 +18,8 @@ import matplotlib.pyplot as plt
 
 def Data_Generator():
     list_test = []
-    for j in os.listdir('./test/'):
-        list_test.append('./test/'+j)
+    for j in os.listdir('./'):
+        list_test.append('./'+j)
     fname = list_test[0]
     return fname