diff --git a/main/experiments/scripts/november_2018/end_to_end_with_2_layers_only_dense_with_augment/deepstrom_classif_end_to_end.py b/main/experiments/scripts/november_2018/end_to_end_with_2_layers_only_dense_with_augment/deepstrom_classif_end_to_end.py
index 60f8e3f602b1a721f35ce56ed409572a2aa2fbf2..9fe84027a6f3fa4bfdfa5c43ccbb222a9b929458 100644
--- a/main/experiments/scripts/november_2018/end_to_end_with_2_layers_only_dense_with_augment/deepstrom_classif_end_to_end.py
+++ b/main/experiments/scripts/november_2018/end_to_end_with_2_layers_only_dense_with_augment/deepstrom_classif_end_to_end.py
@@ -54,13 +54,13 @@ Kernel related:
 
 """
 
-
 import skluc.main.data.mldatasets as dataset
 import numpy as np
 import tensorflow as tf
 from tensorflow.python.keras.layers import Dense
 from tensorflow.python.keras.regularizers import l2
 from tensorflow.python.keras.initializers import he_normal
+from tensorflow.python.keras.preprocessing.image import ImageDataGenerator
 
 from skluc.main.tensorflow_.kernel_approximation.fastfood_layer import FastFoodLayer
 from skluc.main.tensorflow_.kernel_approximation.nystrom_layer import DeepstromLayerEndToEnd
@@ -168,6 +168,13 @@ def main(paraman, resman, printman):
     X_test, y_test = data.test.data, data.test.labels
     X_val, y_val = data.validation.data, data.validation.labels
 
+    datagen = ImageDataGenerator(
+        rotation_range=20,
+        width_shift_range=0.2,
+        height_shift_range=0.2,
+        horizontal_flip=True)
+    datagen.fit(X_train)
+
     paraman.init_kernel_dict(X_train)
 
     # # Model definition
@@ -185,11 +192,8 @@ def main(paraman, resman, printman):
     repr_sub = convnet_model(subs)
 
 
-
     logger.debug(paraman["kernel_dict"])
 
-    input_classifier = None
-
     if paraman["network"] == "deepstrom":
         deepstrom_layer = DeepstromLayerEndToEnd(subsample_size=paraman["--nys-size"],
                                                  kernel_name=paraman["kernel"],
@@ -258,8 +262,8 @@ def main(paraman, resman, printman):
         j = 0
         for i in range(paraman["--num-epoch"]):
             logger.debug(memory_usage())
-
-            for X_batch, Y_batch in batch_generator(X_train, y_train, paraman["--batch-size"], False):
+            k = 0
+            for X_batch, Y_batch in datagen.flow(X_train, y_train, batch_size=paraman["--batch-size"]):
                 if paraman["network"] == "deepstrom":
                     feed_dict = {x: X_batch, y: Y_batch, subs: nys_subsample}
                 else:
@@ -267,9 +271,8 @@ def main(paraman, resman, printman):
                 _, loss, acc, summary_str = sess.run([train_optimizer, cross_entropy, accuracy_op, merged_summary], feed_dict=feed_dict)
                 if j % 100 == 0:
                     logger.info(
-                        "epoch: {}/{}; batch: {}/{}; batch_shape: {}; loss: {}; acc: {}".format(i, paraman["--num-epoch"], j + 1,
-                                                                                                int(data.train[0].shape[
-                                                                                                        0] / paraman["--batch-size"]) + 1,
+                        "epoch: {}/{}; batch: {}/{}; batch_shape: {}; loss: {}; acc: {}".format(i, paraman["--num-epoch"],
+                                                                                                j + 1, int(data.train[0].shape[0] / paraman["--batch-size"]) + 1,
                                                                                                 X_batch.shape, loss,
                                                                                                 acc))
                     if paraman["--tensorboard"]: