diff --git a/Classification_BatchDataset.py b/Classification_BatchDataset.py
index 6c157e6..b1c8474 100644
--- a/Classification_BatchDataset.py
+++ b/Classification_BatchDataset.py
@@ -3,7 +3,7 @@
 """
 from past.builtins import xrange
 import numpy as np
-import scipy.misc as misc
+import imageio
 import pandas as pa
 import re
 import os
@@ -121,7 +121,7 @@ def load_class(self, folder, class_index):
         return None
 
     def load_image(self,folder,image, class_index):
-        image = misc.imread(self.path + "/" + folder + "/" + image)
+        image = imageio.imread(self.path + "/" + folder + "/" + image)
         nr_y = image.shape[0] // self.tile_size[0]
         nr_x = image.shape[1] // self.tile_size[1]
 
diff --git a/Example_Classification.py b/Example_Classification.py
index 2d7326a..15456c3 100644
--- a/Example_Classification.py
+++ b/Example_Classification.py
@@ -103,7 +103,10 @@ def main(unused_argv):
 
     with tf.Session() as sess:
         sess.run(tf.global_variables_initializer())
-        for i in range(5):
+
+        print("Training the network with %d training steps..." % FLAGS.num_steps)
+
+        for i in range(FLAGS.num_steps):
             batch = data_reader.next_batch(FLAGS.batch_size)
             if i % 1000 == 0:
                 train_accuracy = accuracy.eval(feed_dict={
@@ -139,8 +142,9 @@ def main(unused_argv):
                       help='Directory for storing input data')
   parser.add_argument("--batch_size", type=int, default=2, help="batch size for training")
   parser.add_argument("--test_batch_size", type=int, default=200, help="batch size for training")
-  parser.add_argument("--model_path", type=str, default="Models/deepscores_class.ckpt",
+  parser.add_argument("--num_steps", type=int, default=5, help="Number of training steps")
+  parser.add_argument("--model_path", type=str, default="./Models/deepscores_class.ckpt",
                       help="where to store the trained model")
 
   FLAGS, unparsed = parser.parse_known_args()
-  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)
\ No newline at end of file
+  tf.app.run(main=main, argv=[sys.argv[0]] + unparsed)