提交 db6f4cf2 authored 作者: Frederic's avatar Frederic

fix crash for the mean op when the axis is None.

Also assert that the output dtype if float64. This is not important as the Mean op is not used by default.
上级 edbe833a
......@@ -3222,7 +3222,14 @@ class Mean(elemwise.CAReduce):
def perform(self, node, inp, out):
input, = inp
output, = out
output[0] = numpy.mean(input, axis=self.axis[0])
if self.axis is None:
axis = None
else:
axis = self.axis[0]
# numpy.asarray is needed as otherwise we can end up with a
# numpy scalar.
output[0] = numpy.asarray(numpy.mean(input, dtype='float64',
axis=axis))
def c_code(self, node, name, inames, onames, sub):
if self.axis is not None:
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论