onnx cannot load int32 when top10k:1 is the graph output node

Author: simpeng

tf2onnx/graph.py

@@ -300,6 +300,10 @@ def get_dtype(self, name):
         """Get dtype for node."""
         return self._dtypes.get(name)
 
+    def set_dtype(self, name, val):
+        """Set dtype for node."""
+        self._dtypes[name] = val
+
     def get_shape(self, name):
         """Get shape for node."""
         assert isinstance(name, str)

tf2onnx/tfonnx.py

@@ -884,6 +884,9 @@ def topk_op(ctx, node, name, args):
     node.set_attr("k", k[0])
     node.type = "TopK"
     ctx.remove_input(node, node.input[1])
+
+    # the second of TopK operator must be INT64 per ONNX requires. 
+    ctx.set_dtype(name + ":1", onnx_pb.TensorProto.INT64)
     return node