import torch


def get_batch(data_iterator, args):
    """
    Generate a batch func.
    input schema:
        data_iterator: data iterator that implements the torch.utils.data.DataLoader interface
        args: user defined arguments dictionary
    output schema：
        dictionary (python dict()): a dictionary that contains all data used in the model forward step
    """

 
    device = torch.cuda.current_device()

    feat, adj, labels, nodes,traintype = next(data_iterator)

    nodes = nodes.to(device)
    labels = labels[0].to(device)
    feat = feat[0].to(device)#.to_dense()
    adj = adj[0].to(device)
    return feat, adj, labels, nodes, traintype[0]