Fix batching during generation

        with numpy_seed(seed):

            batches = uneven_batches_by_size(

                dataset.src, dataset.dst, max_tokens=max_tokens,

                max_sentences=max_sentences, max_positions=max_positions)

                max_sentences=max_sentences, max_positions=max_positions,

                # FP16: during training keep the batch size a multiple of 8

                required_batch_size_multiple=8,

            )

            frozen_batches = tuple(batches)  # freeze

        def dataloader(b):

def _make_batches(src, dst, indices, max_tokens, max_sentences, max_positions,

                  ignore_invalid_inputs=False, allow_different_src_lens=False):

                  ignore_invalid_inputs=False, allow_different_src_lens=False,

                  required_batch_size_multiple=1):

    batch = []

    mult = required_batch_size_multiple

    def yield_batch(next_idx, num_tokens):

        if len(batch) == 0:

        return False

    sample_len = 0

    sample_lens = []

    ignored = []

    for idx in map(int, indices):

        src_size = src.sizes[idx]

                " Skip this example with --skip-invalid-size-inputs-valid-test"

            ).format(idx, src_size, dst_size, max_positions))

        sample_len = max(sample_len, src_size, dst_size)

        sample_lens.append(max(src_size, dst_size))

        sample_len = max(sample_len, sample_lens[-1])

        num_tokens = (len(batch) + 1) * sample_len

        while yield_batch(idx, num_tokens):

            mod8_len = max(8 * (len(batch) // 8), len(batch) % 8)

        if yield_batch(idx, num_tokens):

            mod8_len = max(mult * (len(batch) // mult), len(batch) % mult)

            yield batch[:mod8_len]

            batch = batch[mod8_len:]

            sample_len = max([max(src.sizes[id], dst.sizes[id]) for id in batch]) if len(batch) > 0 else 0

            sample_len = max(sample_len, src_size, dst_size)

            num_tokens = (len(batch) + 1) * sample_len

            sample_lens = sample_lens[mod8_len:]

            sample_len = max(sample_lens) if len(sample_lens) > 0 else 0

        batch.append(idx)

def batches_by_size(src, dst, max_tokens=None, max_sentences=None,

                    max_positions=(1024, 1024), ignore_invalid_inputs=False,

                    descending=False):

                    descending=False, required_batch_size_multiple=1):

    """Returns batches of indices sorted by size. Sequences with different

    source lengths are not allowed in the same batch."""

    assert isinstance(src, IndexedDataset) and (dst is None or isinstance(dst, IndexedDataset))

        indices = np.flip(indices, 0)

    return list(_make_batches(

        src, dst, indices, max_tokens, max_sentences, max_positions,

        ignore_invalid_inputs, allow_different_src_lens=False))

        ignore_invalid_inputs, allow_different_src_lens=False,

        required_batch_size_multiple=required_batch_size_multiple,

    ))

def uneven_batches_by_size(src, dst, max_tokens=None, max_sentences=None, max_positions=(1024, 1024)):

def uneven_batches_by_size(src, dst, max_tokens=None, max_sentences=None,

                           max_positions=(1024, 1024),

                           required_batch_size_multiple=1):

    """Returns batches of indices bucketed by size. Batches may contain

    sequences of different lengths."""

    assert isinstance(src, IndexedDataset) and isinstance(dst, IndexedDataset)

    batches = list(_make_batches(

        src, dst, indices, max_tokens, max_sentences, max_positions,

        ignore_invalid_inputs=True, allow_different_src_lens=True))

        ignore_invalid_inputs=True, allow_different_src_lens=True,

        required_batch_size_multiple=required_batch_size_multiple,

    ))

    return batches