Fix a bug

2020-10-20 14:30:00 +08:00 · 2020-10-20 14:30:00 +08:00 · 44d89ae217
parent c44d32fa27
commit 44d89ae217
1 changed files with 4 additions and 1 deletions
--- a/tools/infer/infer.py
+++ b/tools/infer/infer.py
@ -53,6 +53,9 @@ def create_predictor(args):
            out = fluid.layers.softmax(out)
        return out

+    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()