import tensorflow.contrib.eager as tfe

                if policy['type'] == 'gcnn':
                    # load model
                    sys.path.insert(0, os.path.abspath(f"models/{policy['name']}"))
                    import model
                    importlib.reload(model)
                    del sys.path[0]
                    policy['model'] = model.GCNPolicy()
                    policy['model'].restore_state(f"trained_models/{args.problem}/{policy['name']}/{seed}/best_params.pkl")
                    policy['model'].call = tfe.defun(policy['model'].call, input_signature=policy['model'].input_signature)
                    policy['batch_datatypes'] = [tf.float32, tf.int32, tf.float32,
                            tf.float32, tf.int32, tf.int32, tf.int32, tf.int32, tf.int32, tf.float32]
                    policy['batch_fun'] = load_batch_gcnn
                else:
                    # load feature normalization parameters
                    try:
                        with open(f"trained_models/{args.problem}/{policy['name']}/{seed}/normalization.pkl", 'rb') as f:
