Skip to content

Commit ef8cb8f

Browse files
Merge pull request #7816 from wanghaoshuang/infer_prog
Make get_inference_program support for Evaluator.
2 parents 32a5dfd + a6a79c3 commit ef8cb8f

File tree

1 file changed

+9
-2
lines changed
  • python/paddle/v2/fluid

1 file changed

+9
-2
lines changed

python/paddle/v2/fluid/io.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
import os
1616
import cPickle as pickle
1717

18+
from paddle.v2.fluid.evaluator import Evaluator
1819
from paddle.v2.fluid.framework import Program, Parameter, default_main_program, Variable
1920
from . import core
2021

@@ -187,8 +188,14 @@ def get_inference_program(target_vars, main_program=None):
187188
main_program = default_main_program()
188189
if not isinstance(target_vars, list):
189190
target_vars = [target_vars]
190-
191-
pruned_program = main_program.prune(targets=target_vars)
191+
vars = []
192+
for var in target_vars:
193+
if isinstance(var, Evaluator):
194+
vars.append(var.states)
195+
vars.append(var.metrics)
196+
else:
197+
vars.append(var)
198+
pruned_program = main_program.prune(targets=vars)
192199
inference_program = pruned_program.inference_optimize()
193200
return inference_program
194201

0 commit comments

Comments
 (0)