From cbb9f9f31775b61666202cece2ce66d3a923a185 Mon Sep 17 00:00:00 2001 From: Lai Wei Date: Mon, 12 Nov 2018 17:05:00 -0800 Subject: [PATCH] fix train mnist for inception-bn and resnet --- example/image-classification/train_mnist.py | 1 + 1 file changed, 1 insertion(+) diff --git a/example/image-classification/train_mnist.py b/example/image-classification/train_mnist.py index 2bc4289318d9..17a5a372c0c2 100644 --- 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(