code: def batch(input_iter, batch_size=32):
    
    input_iter = iter(input_iter)
    next_ = list(itertools.islice(input_iter, batch_size))
    while next_:
      yield next_
      next_ = list(itertools.islice(input_iter, batch_size))
docs: Batches data from an iterator that returns single items at a time.