Make sure the final decoded ids shape is [beam_width, batch_size, max_time] when using beam search, and [batch_size, max_time] when using basic decoder.
Make sure the final decoded ids shape is [beam_width, batch_size, max_time] when
using beam search, and [batch_size, max_time] when using basic decoder.

Fixes #180

PiperOrigin-RevId: 176599755

tensorflow/nmt,tensorflow/nmt