Fix the Bert2Bert call when padded_decode is set to True

[LoicDagnas]

Description

As mentioned in the following issue, it was impossible to call a Bert2Bert model with the configuration option padded_decode set to True.

It comes from how the beam search modules deals with batch size. The special case implemented when padded_decode is True doesn't seem useful and caused the following error:

TypeError: Failed to convert object of type <class 'list'> to Tensor. Contents: [None, 1]. Consider casting elements to a supported type.

Type of change

• [x] Bug fix (non-breaking change which fixes an issue)

Tests

I have added a naive unittest which failed before my modification and passes now.

Checklist

• [x] I have signed the Contributor License Agreement.
• [x] I have read guidelines for pull request.
• [x] My code follows the coding guidelines.
• [x] I have performed a self code review of my own code.
• [ ] My changes generate no new warnings.
• [x] I have added tests that prove my fix is effective or that my feature works.

[LoicDagnas]

@saberkun or anyone else if you have a bit of time for this one, it will be very helpful. It is quite critical to make my model serveable in ONNX run time or in TPU vms.

If I have to modify the code or add unittest, just tell me :)

Thanks in advance !

[laxmareddyp]

Hi @LoicDagnas,

In future, please open an issue and discuss with us before submitting a pull request. If the change does not align with our best practices, we can provide early feedback. Please review guidelines here: https://github.com/tensorflow/models/wiki/Submitting-a-pull-request

Thanks.