fix onnx shape dtype (#4528)

093405d8 · masahi · Tianqi Chen · 8e3b5d39 · 093405d8
Commit 093405d8 authored Dec 17, 2019 by masahi Committed by Tianqi Chen Dec 16, 2019
Hide whitespace changes
Inline Side-by-side

Showing with 1 additions and 2 deletions

python/tvm/relay/frontend/onnx.py
+1 -2

No files found.
--- 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.