From 79a3dbc3e08b026a61eb3c324d6a7c942bc66323 Mon Sep 17 00:00:00 2001 From: littletomatodonkey <2120160898@bit.edu.cn> Date: Sun, 6 Dec 2020 17:52:55 +0800 Subject: [PATCH] fix export effnet model (#451) --- tools/export_model.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/export_model.py b/tools/export_model.py index d911856d..b3d2fe2f 100644 --- a/tools/export_model.py +++ b/tools/export_model.py @@ -55,7 +55,10 @@ def create_model(args, model, input, class_dim=1000): def main(): args = parse_args() - model = architectures.__dict__[args.model]() + if "EfficientNet" in args.model: + model = architectures.__dict__[args.model](is_test=True) + else: + model = architectures.__dict__[args.model]() place = fluid.CPUPlace() exe = fluid.Executor(place) -- GitLab