fix train mnist for inception-bn and resnet (#13239)

apache · Nov 15, 2018 · 1ef83c9 · 1ef83c9
1 parent 339e085
commit 1ef83c9
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/example/image-classification/train_mnist.py b/example/image-classification/train_mnist.py
@@ -72,6 +72,7 @@ def get_mnist_iter(args, kv):
                         help='the number of training examples')
 
     parser.add_argument('--add_stn',  action="store_true", default=False, help='Add Spatial Transformer Network Layer (lenet only)')
+    parser.add_argument('--image_shape', default='1, 28, 28', help='shape of training images')
 
     fit.add_fit_args(parser)
     parser.set_defaults(