def METHOD_NAME(model, data_kwargs: dict, pp_rank: int, chunk: int, stage_num: int):
    torch.manual_seed(1024)
    METHOD_NAME = create_partition_module(pp_rank, stage_num, model, data_kwargs)
    return METHOD_NAME