code runs

                                group=mpu.get_model_parallel_group())

    context_length = context_length_tensor.min().item()

    tokens, attention_mask, position_ids = get_batch(context_tokens_tensor, args)

    tokens, attention_mask, position_ids = get_batch(context_tokens_tensor)

    batch_token_iterator = sample_sequence_batch(model, context_tokens_tensor,

                                                 context_length_tensor,

    # Generate samples.

    if args.num_samples == 0:

        assert args.batch_size == 1

        args.batch_size = 1

        if args.sample_input_file != "":

            generate_samples_input_from_file(model)

        else:

    # Checks.

    assert args.hidden_size % args.num_attention_heads == 0

    if args.seq_length is not None:

        assert args.max_position_embeddings >= args.seq_length

    if args.lr is not None:

        assert args.min_lr <= args.lr

    if args.save is not None:

        assert args.save_interval is not None

def _add_training_args(parser):

    group = parser.add_argument_group(title='training')

    group.add_argument('--batch-size', type=int, required=True,

    group.add_argument('--batch-size', type=int, default=None,

                       help='Batch size per model instance (local batch size). '

                       'Global batch size is local batch size times data '

                       'parallel size.')

                       help='Path to the vocab file.')

    group.add_argument('--merge-file', type=str, default=None,

                       help='Path to the BPE merge file.')

    group.add_argument('--seq-length', type=int, required=True,

    group.add_argument('--seq-length', type=int, default=None,

                       help="Maximum sequence length to process.")

    group.add_argument('--mask-prob', type=float, default=0.15,

                       help='Probability of replacing a token with mask.')

def add_data_args_(parser):

    """Train/valid/test data arguments."""