INNER CODE UNIT · Python

batchify

IBM/transition-amr-parser · src/ibm_neural_aligner/main.py:573

def batchify(items, cuda=False, train=False):

    if cuda and torch.cuda.is_available():
        device = torch.cuda.current_device()
    else:
        device = None

    dtypes = {
        'text_tokens': torch.long,
        'amr_tokens': torch.long,
        'amr_node_ids': torch.long,
        'amr_node_mask': torch.bool,
    }

    batch_map = {}
    for k in dtypes.keys():
        if k not in items[0]:
            continue

View source record →