Yields (labels, np_image) tuples.

def _load_data(path, labels_number=1):
  """Yields (labels, np_image) tuples."""
  with tf.io.gfile.GFile(path, "rb") as f:
    data = f.read()
  offset = 0
  max_offset = len(data) - 1
  while offset < max_offset:
    labels = np.frombuffer(data, dtype=np.uint8, count=labels_number,
                           offset=offset).reshape((labels_number,))
    # 1 byte per label, 1024 * 3 = 3072 bytes for the image.
    offset += labels_number
    img = (np.frombuffer(data, dtype=np.uint8, count=3072, offset=offset)
           .reshape((3, _CIFAR_IMAGE_SIZE, _CIFAR_IMAGE_SIZE))
           .transpose((1, 2, 0))
          )
    offset += 3072
    yield labels, img