fix onnx shape dtype (apache#4528)

neo-ai · Jan 11, 2020 · bd925da · bd925da
1 parent 27674bd
commit bd925da
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py
@@ -687,8 +687,7 @@ class Shape(OnnxOpConverter):
 
     @classmethod
     def _impl_v1(cls, inputs, attr, params):
-        # TODO(@jroesch): use shape_of once it has been fixed)
-        return _op.shape_of(inputs[0])
+        return _op.shape_of(inputs[0], "int64")
 
 class Cast(OnnxOpConverter):
     """ Operator converter for Cast.