diff --git a/nnvm/python/nnvm/compiler/build_module.py b/nnvm/python/nnvm/compiler/build_module.py index 3080baf1d79dc8f7485f398a94cee225d21567f0..86fa08ec1bfcf15db1bd85804072c0ba93032070 100644 --- a/nnvm/python/nnvm/compiler/build_module.py +++ b/nnvm/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)