From 6fcef38122b34acd19c226f5938c938f256a25c4 Mon Sep 17 00:00:00 2001 From: SunAhong1993 <48579383+SunAhong1993@users.noreply.github.com> Date: Thu, 28 Mar 2019 16:33:59 +0800 Subject: [PATCH] Update convert.py --- caffe2fluid/convert.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/caffe2fluid/convert.py b/caffe2fluid/convert.py index 95bafb8..adce7d0 100755 --- a/caffe2fluid/convert.py +++ b/caffe2fluid/convert.py @@ -41,8 +41,12 @@ def convert(def_path, caffemodel_path, data_output_path, code_output_path, np.save(data_out, data) if code_output_path: print_stderr('Saving source...') + s = sys.version with open(code_output_path, 'wb') as src_out: - src_out.write(str.encode(transformer.transform_source())) + if s.startswith('2'): + src_out.write(transformer.transform_source()) + else: + src_out.write(str.encode(transformer.transform_source())) print_stderr('set env variable before using converted model '\ 'if used custom_layers:') custom_pk_path = os.path.dirname(os.path.abspath(__file__)) -- GitLab