raise error when cannot get shape in reshape (dmlc#467)

* raise error when cannot get shape in reshape * fix pylint
lebeg · May 5, 2018 · 8c0a103 · 8c0a103
1 parent 63bd093
commit 8c0a103
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/python/nnvm/frontend/onnx.py b/python/nnvm/frontend/onnx.py
@@ -246,16 +246,20 @@ def _impl_v1(cls, inputs, attr, params):
 
 
 class Reshape(OnnxOpConverter):
+    """ Operator converter for Reshape.
+    """
 
     @classmethod
     def _impl_v1(cls, inputs, attr, params):
         return _sym.reshape(inputs[0], shape=attr['shape'])
 
     @classmethod
     def _impl_v5(cls, inputs, attr, params):
-        return _sym.reshape(
-            inputs[0],
-            shape=tuple(params[inputs[1].list_output_names()[0]].asnumpy()))
+        if inputs[1].list_output_names()[0] in params:
+            shape = tuple(params[inputs[1].list_output_names()[0]].asnumpy())
+        else:
+            raise RuntimeError('Shape is not contained in graph initializer.')
+        return _sym.reshape(inputs[0], shape=shape)
 
 
 class Scale(OnnxOpConverter):