def Prefetch(n_prefetch=2):  # pylint: disable=invalid-name
  """Pre-fetches a number of examples from generator in a separate process."""
  def prefetch(generator):
    in_q, out_q = mp.Queue(), mp.Queue()
    p = mp.Process(target=_generator_process, args=(generator, in_q, out_q))
    for _ in range(n_prefetch):
      in_q.put(None)
    p.start()
    while True:
      yield out_q.get()
      in_q.put(None)
  return prefetch
Pre-fetches a number of examples from generator in a separate process.
trax/data/inputs.py