diff --git a/caffe2fluid/examples/imagenet/infer.py b/caffe2fluid/examples/imagenet/infer.py
index 6fca82bba82bd34c9be74850a41dfa57fccc7f8a..3c5c47457180972d8720f6c279ae99eb5b1fc9e8 100644
--- a/caffe2fluid/examples/imagenet/infer.py
+++ b/caffe2fluid/examples/imagenet/infer.py
@@ -2,7 +2,7 @@
 
 #function:
 #   a demo to show how to use the converted model genereated by caffe2fluid
-#   
+#
 #notes:
 #   only support imagenet data
 
@@ -59,8 +59,9 @@ def build_model(net_file, net_name):
     inputs_dict = MyNet.input_shapes()
     input_name = list(inputs_dict.keys())[0]
     input_shape = inputs_dict[input_name]
-    images = fluid.layers.data(
-        name=input_name, shape=input_shape, dtype='float32')
+    images = fluid.layers.data(name=input_name,
+                               shape=input_shape,
+                               dtype='float32')
     #label = fluid.layers.data(name='label', shape=[1], dtype='int64')
 
     net = MyNet({input_name: images})
@@ -172,7 +173,7 @@ def load_model(exe, place, net_file, net_name, net_weight, debug):
 
 def get_shape(fluid, program, name):
     for var in program.list_vars():
-        if var.type == 'Input':
+        if var.name == name:
             return list(var.shape[1:])
 
     raise ValueError('not found shape for input layer[%s], '
diff --git a/caffe2fluid/examples/imagenet/tools/diff.sh b/caffe2fluid/examples/imagenet/tools/diff.sh
index 8a561a821f8d10d2e8537d83908029532d3cde15..eca3f79e9fee4c116401d2a6db321d7c6a5c2c76 100755
--- a/caffe2fluid/examples/imagenet/tools/diff.sh
+++ b/caffe2fluid/examples/imagenet/tools/diff.sh
@@ -49,7 +49,7 @@ if [[ -z $PYTHON ]];then
     PYTHON=`which python`
 fi
 $PYTHON ../../convert.py \
-        --npy_path $proto_file \
+        --def_path $proto_file \
         --caffemodel $caffemodel_file \
         --data-output-path $weight_file\
         --code-output-path $net_file