Generate Token in Yuan2.0

[hungptit123]

I see the code generate token. But I want to generate more tokens (2, 3, 4, 5 tokens) in one inference. But the code can do it.
You can help me generate more tokens in one inference.

The step = 1 in loop

for context_length in range(prompt_length, final_sequence_length, 1):

        # Pick the slice that we need to pass through the network.
        tokens2use = tokens[:, prev_context_length:context_length]
        positions2use = position_ids[:, prev_context_length:context_length]
        attention_mask2use = attention_mask[
            ..., prev_context_length:context_length, :context_length]

        # logits will be meanigful only in the last pipeline stage.
        logits = forward_step(tokens2use, positions2use, attention_mask2use)

-> That OK
But step = 2 in loop , this is error.

for context_length in range(prompt_length, final_sequence_length, 2):

        # Pick the slice that we need to pass through the network.
        tokens2use = tokens[:, prev_context_length:context_length]
        positions2use = position_ids[:, prev_context_length:context_length]
        attention_mask2use = attention_mask[
            ..., prev_context_length:context_length, :context_length]

        # logits will be meanigful only in the last pipeline stage.
        logits = forward_step(tokens2use, positions2use, attention_mask2use)

[Shawn-IEITSystems]

@zhaoxudong01-ieisystem

[zhaoxudong01]

Can you give more error message or full log file.

[Shawn-IEITSystems]

@hungptit123 Have you fixed this issue?