File tree Expand file tree Collapse file tree 1 file changed +9
-2
lines changed Expand file tree Collapse file tree 1 file changed +9
-2
lines changed Original file line number Diff line number Diff line change 15
15
import os
16
16
import cPickle as pickle
17
17
18
+ from paddle .v2 .fluid .evaluator import Evaluator
18
19
from paddle .v2 .fluid .framework import Program , Parameter , default_main_program , Variable
19
20
from . import core
20
21
@@ -187,8 +188,14 @@ def get_inference_program(target_vars, main_program=None):
187
188
main_program = default_main_program ()
188
189
if not isinstance (target_vars , list ):
189
190
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 )
192
199
inference_program = pruned_program .inference_optimize ()
193
200
return inference_program
194
201
You can’t perform that action at this time.
0 commit comments