import tensorflow as tf

        types[rconst.VALID_POINT_MASK] = np.bool
        shapes[rconst.VALID_POINT_MASK] = tf.TensorShape([batch_size])

        types = (types, np.bool)
        shapes = (shapes, tf.TensorShape([batch_size]))

      else:
        types[rconst.DUPLICATE_MASK] = np.bool
        shapes[rconst.DUPLICATE_MASK] = tf.TensorShape([batch_size])

      data_generator = functools.partial(
          self.data_generator, epochs_between_evals=epochs_between_evals)
      dataset = tf.data.Dataset.from_generator(
          generator=data_generator, output_types=types,
          output_shapes=shapes)

