code: def dataset_to_stream(dataset, input_name, num_chunks=0, append_targets=False):
    
    for example in tfds.as_numpy(dataset):
      inp, out = example[0][input_name], example[1]
      if len(out.shape) > 1 and out.shape[-1] == 1:
        out = np.squeeze(out, axis=-1)
      if num_chunks > 0:
        inp = np.split(inp, num_chunks, axis=1)
        out = np.split(out, num_chunks, axis=1)
      if append_targets:
        inp = (inp, out)
      yield inp, out
docs: Takes a tf.Dataset and creates a numpy stream of ready batches.