Small bugfix to make sure refactored code works with interleaved schedule

    """Build the model."""

    args = get_args()

    # Build model on cpu.

    pre_process = mpu.is_pipeline_first_stage()

    post_process = mpu.is_pipeline_last_stage()

    # Build model.

    if mpu.get_pipeline_model_parallel_world_size() > 1 and \

       args.virtual_pipeline_model_parallel_size is not None:

        model = []

        for i in range(args.virtual_pipeline_model_parallel_size):

            mpu.set_virtual_pipeline_model_parallel_rank(i)

            # Set pre_process and post_process only after virtual rank is set.

            pre_process = mpu.is_pipeline_first_stage()

            post_process = mpu.is_pipeline_last_stage()

            this_model = model_provider_func(

                pre_process=pre_process,

                post_process=post_process

            )

            model.append(this_model)

    else:

        pre_process = mpu.is_pipeline_first_stage()

        post_process = mpu.is_pipeline_last_stage()

        model = model_provider_func(

            pre_process=pre_process,

            post_process=post_process