Fix function of graph wrapper. (PaddlePaddle#35)

paddleslim/core/graph_wrapper.py

@@ -96,6 +96,7 @@ def outputs(self):
     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 @@ def var(self, name):
         """
         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):
         """