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,6 +55,9 @@ def create_model(args, model, input, class_dim=1000):
 def main():
    args = parse_args()
+    if "EfficientNet" in args.model:
+        model = architectures.__dict__[args.model](is_test=True)
+    else:
        model = architectures.__dict__[args.model]()
    place = fluid.CPUPlace()