def METHOD_NAME(reader, dev_count):
    """multi device"""
    if dev_count == 1:
        for batch in reader:
            yield batch
    else:
        batches = []
        for batch in reader:
            batches.append(batch)
            if len(batches) == dev_count:
                yield batches
                batches = []