From 026e88a340f3f89deb27e1d8060b9fb53db32c90 Mon Sep 17 00:00:00 2001
From: wanghaoshuang <wanghaoshuang@baidu.com>
Date: Wed, 24 Jan 2018 14:34:54 +0800
Subject: [PATCH] Make get_inference_program support for Evaluator.
---
python/paddle/v2/fluid/io.py | 3 +++
1 file changed, 3 insertions(+)
diff --git a/python/paddle/v2/fluid/io.py b/python/paddle/v2/fluid/io.py
index 499df05e59..d440fab9a5 100644
--- a/python/paddle/v2/fluid/io.py
+++ b/python/paddle/v2/fluid/io.py
@@ -14,6 +14,7 @@
import os
import cPickle as pickle
+from paddle.v2.fluid.evaluator import Evaluator
from paddle.v2.fluid.framework import Program, Parameter, default_main_program, Variable
__all__ = [
@@ -183,6 +184,8 @@ def load_persistables(executor, dirname, main_program=None):
def get_inference_program(target_vars, main_program=None):
if main_program is None:
main_program = default_main_program()
+ if isinstance(target_vars, Evaluator):
+ target_vars = target_vars.states + target_vars.metrics
if not isinstance(target_vars, list):
target_vars = [target_vars]