From d7e074c265328a4ebe02048905349878977c43f0 Mon Sep 17 00:00:00 2001
From: whs <wanghaoshuang@baidu.com>
Date: Tue, 14 Jan 2020 17:43:59 +0800
Subject: [PATCH] Fix function of graph wrapper. (#35)

---
 paddleslim/core/graph_wrapper.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/paddleslim/core/graph_wrapper.py b/paddleslim/core/graph_wrapper.py
index 7ed99d20..1feedd12 100644
--- a/paddleslim/core/graph_wrapper.py
+++ b/paddleslim/core/graph_wrapper.py
@@ -96,6 +96,7 @@ class VarWrapper(object):
     def is_parameter(self):
         return isinstance(self._var, Parameter)
 
+
 class OpWrapper(object):
     def __init__(self, op, graph):
         assert isinstance(graph, GraphWrapper)
@@ -270,7 +271,10 @@ class GraphWrapper(object):
         """
         Get the variable by variable name.
         """
-        return VarWrapper(self.program.global_block().var(name), self)
+        for block in self.program.blocks:
+            if block.has_var(name):
+                return VarWrapper(block.var(name), self)
+        return None
 
     def clone(self, for_test=False):
         """
-- 
GitLab