def METHOD_NAME(args):
    if not args.use_cuda:
        paddle.set_device("cpu")
    if paddle.distributed.get_world_size() > 1:
        paddle.distributed.init_parallel_env()
    if args.edge_file:
        graph = load_from_file(args.edge_file)
    else:
        graph = load(args.dataset)
    model = SkipGramModel(
        graph.num_nodes,
        args.embed_size,
        args.neg_num,
        sparse=not args.use_cuda,
        shared_embedding=args.shared_embedding)
    model = paddle.DataParallel(model)
    train_ds = ShardedDataset(graph.nodes, repeat=args.epoch)
    train_steps = int(len(train_ds) // args.batch_size)
    log.info("train_steps: %s" % train_steps)
    scheduler = paddle.optimizer.lr.PolynomialDecay(
        learning_rate=args.learning_rate,
        decay_steps=train_steps,
        end_lr=0.0001)
    optim = Adam(learning_rate=scheduler, parameters=model.parameters())
    collate_fn = BatchRandWalk(graph, args.walk_len, args.win_size,
                               args.neg_num, args.neg_sample_type)
    data_loader = Dataloader(
        train_ds,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.sample_workers,
        collate_fn=collate_fn)
    train_loss = train(model, data_loader, optim)
    paddle.save(model.state_dict(), "model.pdparams")