Fix a bug

44d89ae2 · gaotingquan · c44d32fa · 44d89ae2
隐藏空白更改
内联并排

Showing with 4 addition and 1 deletion

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

未找到文件。
--- a/tools/infer/infer.py
+++ b/tools/infer/infer.py
@@ -53,7 +53,10 @@ def create_predictor(args):
            out = fluid.layers.softmax(out)
        return out

-    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.CUDAPlace(0) if args.use_gpu else fluid.CPUPlace()
    exe = fluid.Executor(place)