import tensorflow as tf

                    if policy_name.startswith('svmrank'):
                        policy['model'] = svmrank.Model().read(f"trained_models/{args.problem}/{policy['name']}/{seed}/model.txt")
                    else:
                        with open(f"trained_models/{args.problem}/{policy['name']}/{seed}/model.pkl", 'rb') as f:
                            policy['model'] = pickle.load(f)

                    # load feature specifications
                    with open(f"trained_models/{args.problem}/{policy['name']}/{seed}/feat_specs.pkl", 'rb') as f:
                        feat_specs = pickle.load(f)

                    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']))
                test_data = test_data.prefetch(2)

                test_kacc = process(policy, test_data, top_k)
                print(f"  {seed} " + " ".join([f"acc@{k}: {100*acc:4.1f}" for k, acc in zip(top_k, test_kacc)]))

                writer.writerow({
                    **{
                        'policy': f"{policy['type']}:{policy['name']}",
                        'seed': seed,
                    },
                    **{
                        f'acc@{k}': test_kacc[i] for i, k in enumerate(top_k)
                    },
