diff --git a/main.py b/main.py
index 1305231..8cad634 100644
--- a/main.py
+++ b/main.py
@@ -70,7 +70,7 @@ def train(conf, data):
         conf.img_height = 28
         conf.img_width = 28
         conf.channel = 1
-        conf.num_batches = 10#data.train.num_examples // conf.batch_size
+        conf.num_batches = data.train.num_examples // conf.batch_size
     else:
         from keras.datasets import cifar10
         data = cifar10.load_data()