提交 529eb383 authored 作者: Arnaud Bergeron's avatar Arnaud Bergeron

Make sure the axes array will be integer-typed.

上级 a3242677
......@@ -1360,7 +1360,8 @@ class MaxAndArgmax(Op):
dtype=node.outputs[0].dtype)
# Numpy does not support multiple axes for argmax
# Work around
keep_axes = numpy.array([i for i in range(x.ndim) if i not in axes])
keep_axes = numpy.array([i for i in range(x.ndim) if i not in axes],
dtype='int64')
# Not-reduced axes in front
transposed_x = numpy.transpose(x, numpy.concatenate((keep_axes, axes)))
reshaped_x = transposed_x.reshape(transposed_x.shape[:len(keep_axes)] +
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论