def move_batch(batch, device):
    """Moves a batch of dictionary to the corresponding local device"""
    res = {}
    for k in batch.keys():
        if isinstance(batch[k], list):
            res[k] = [i.to(device) for i in batch[k]]
        else:
            res[k] = batch[k].to(device)
    return res