add sanity check to input shape type (#469)

dmlc · May 8, 2018 · 319d9f1 · 319d9f1
1 parent 8c0a103
commit 319d9f1
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/python/nnvm/compiler/build_module.py b/python/nnvm/compiler/build_module.py
@@ -241,6 +241,10 @@ def build(graph, target=None, shape=None, dtype="float32",
     shape = shape if shape else {}
     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")
+
     cfg = BuildConfig.current
     graph = graph if isinstance(graph, _graph.Graph) else _graph.create(graph)
     shape, dtype = _update_shape_dtype(shape, dtype, params)