fix export effnet model (#451)

79a3dbc3 · littletomatodonkey · GitHub · 27013344 · 79a3dbc3
隐藏空白更改
内联并排

Showing with 4 addition and 1 deletion

tools/export_model.py tools/export_model.py +4 -1

未找到文件。
--- 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)