From af9a4d4594c6a986dea6ff89b745a0885293457f Mon Sep 17 00:00:00 2001
From: SunAhong1993 <sunyanfang01@baidu.com>
Date: Tue, 23 Jul 2019 12:08:37 +0800
Subject: [PATCH] fix the bug of infer.py

---
 caffe2fluid/examples/imagenet/infer.py      | 9 +++++----
 caffe2fluid/examples/imagenet/tools/diff.sh | 2 +-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/caffe2fluid/examples/imagenet/infer.py b/caffe2fluid/examples/imagenet/infer.py
index 6fca82b..3c5c474 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 8a561a8..eca3f79 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
-- 
GitLab