def METHOD_NAME(model, data_kwargs, num_microbatches, pp_rank: int, chunk: int, stage_num: int):
    module = create_partition_module(pp_rank, stage_num, model, data_kwargs, num_microbatches)
    return module