def CastTo(dtype=np.int32, indices=(0, 1,), debug=False):  # pylint: disable=invalid-name
  """Casts the given indices to the given dtype."""
  def _cast_fn(generator):
    debug_count = 0
    for example in generator:
      debug_count += 1
      assert isinstance(example, tuple)
      example = list(example)
      dtype_mismatch = False
      original_index_and_dtype = []
      for i in range(len(example)):
        if i not in indices:
          continue
        original_type = example[i].dtype
        if original_type != dtype:
          if not (original_type == np.int64 and dtype == np.int32):
            # Downcasting from np.int64 to np.int32 is OK
            original_index_and_dtype.append((i, original_type))
          example[i] = example[i].astype(dtype)
          dtype_mismatch = True
      if debug and dtype_mismatch and original_index_and_dtype:
        logging.info('dtype mismatch in example[%d] = %r was earlier: %r',
                     debug_count, example, original_index_and_dtype)
      yield tuple(example)
  return _cast_fn
Casts the given indices to the given dtype.
trax/data/inputs.py