Fix --prefix-size

    group.add_argument('--score-reference', action='store_true',

                       help='just score the reference translation')

    group.add_argument('--prefix-size', default=0, type=int, metavar='PS',

                       help=('initialize generation by target prefix of given length'))

                       help='initialize generation by target prefix of given length')

    group.add_argument('--sampling', action='store_true',

                       help='sample hypotheses instead of using beam search')

    return group

                else:

                    # take the best 2 x beam_size predictions. We'll choose the first

                    # beam_size of these which don't predict eos to continue with.

                    torch.topk(

                        probs.view(bsz, -1),

                        k=min(cand_size, probs.view(bsz, -1).size(1) - 1),  # -1 so we never select pad

            assert step < maxlen

            if len(finalized_sents) > 0:

                # construct batch_idxs which holds indices of batches to keep for the next pass

                new_bsz = bsz - len(finalized_sents)

                # construct batch_idxs which holds indices of batches to keep for the next pass

                batch_mask = torch.ones(bsz).type_as(cand_indices)

                batch_mask[cand_indices.new(finalized_sents)] = 0

                batch_idxs = batch_mask.nonzero().squeeze(-1)

                cand_scores = cand_scores[batch_idxs]

                cand_indices = cand_indices[batch_idxs]

                if prefix_tokens is not None:

                    prefix_tokens = prefix_tokens[batch_idxs]

                scores = scores.view(bsz, -1)[batch_idxs].view(new_bsz * beam_size, -1)

                scores_buf.resize_as_(scores)