diff --git a/nnvm/python/nnvm/compiler/build_module.py b/nnvm/python/nnvm/compiler/build_module.py
index 6fab4460b427ed1b0c9cf0d19db397b4f486f1a0..b04d494788301102ed570e52dae0687cdad225c2 100644
--- a/nnvm/python/nnvm/compiler/build_module.py
+++ b/nnvm/python/nnvm/compiler/build_module.py
@@ -251,8 +251,8 @@ def build(graph, target=None, shape=None, dtype="float32",
         if not isinstance(shape, dict):
             raise TypeError("require shape to be dict")
         for value in shape.values():
-            if not all(isinstance(x, int) for x in value):
-                raise TypeError("shape value must be int iterator")
+            if not all(isinstance(x, tvm._ffi.base.integer_types) for x in value):
+                raise TypeError("shape value must be Integer types iterator")
 
         cfg = BuildConfig.current
         graph = graph if isinstance(graph, _graph.Graph) else _graph.create(graph)