Commit 67fe4db0 by lixiaoquan Committed by Tianqi Chen

[NNVM] Fix dtype of output of mean. (#2334)

dtype of count is the same as dtype of inputs[0] when created, but its type may
  change when multiplied by inputs[0]->shape[i]. Which causes dtype of
  output is not same as dtype of input.
parent b076cad5
......@@ -352,7 +352,7 @@ Example::
Expr count = make_const(inputs[0]->dtype, 1);
for (auto& i : r_axes) {
count *= inputs[0]->shape[i];
count *= cast(inputs[0]->dtype, inputs[0]->shape[i]);
}
return Array<Tensor>{
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment