def input_fn_wrapper(*args, **kwargs):
    dataset = input_fn(self._bridge, self._trainer_master)

    def mapper(features, *args):
        features.update(self._preprocess_fids(features.pop('fids'), slot_configs))
        return (((features,) + args) if args else features)
    dataset = dataset.map(mapper, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.prefetch(2)
    return dataset
def input_fn_wrapper(*args, **kwargs):
    dataset = input_fn(self._bridge, self._trainer_master)

    def mapper(features, *args):
        features.update(self._preprocess_fids(features.pop('fids'), slot_configs))
        return (((features,) + args) if args else features)
    dataset = dataset.map(mapper, num_parallel_calls=tf.data.experimental.AUTOTUNE)
    dataset = dataset.prefetch(2)
    return dataset