Fix bug: the output shape is not right in benchmark.

See merge request !747