import tensorflow as tf


                    policy['batch_datatypes'] = [tf.float32, tf.int32, tf.int32, tf.float32]
                    policy['batch_fun'] = lambda x: load_batch_flat(x, feat_specs['type'], feat_specs['augment'], feat_specs['qbnorm'])

                test_data = tf.data.Dataset.from_tensor_slices(test_files)
                test_data = test_data.batch(test_batch_size)
                test_data = test_data.map(lambda x: tf.py_func(
                    policy['batch_fun'], [x], policy['batch_datatypes']))
