def get_model_fn():

    def model_fn(features, labels, mems, is_training):
        (total_loss, new_mems, monitor_dict) = function_builder.get_loss(FLAGS, features, labels, mems, is_training)
        num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
        tf.logging.info('#params: {}'.format(num_params))
        assert is_training
        all_vars = tf.trainable_variables()
        grads = tf.gradients(total_loss, all_vars)
        grads_and_vars = list(zip(grads, all_vars))
        return (total_loss, new_mems, grads_and_vars)
    return model_fn
def get_model_fn():

    def model_fn(features, labels, mems, is_training):
        (total_loss, new_mems, monitor_dict) = function_builder.get_loss(FLAGS, features, labels, mems, is_training)
        num_params = sum([np.prod(v.shape) for v in tf.trainable_variables()])
        tf.logging.info('#params: {}'.format(num_params))
        assert is_training
        all_vars = tf.trainable_variables()
        grads = tf.gradients(total_loss, all_vars)
        grads_and_vars = list(zip(grads, all_vars))
        return (total_loss, new_mems, grads_and_vars)
    return model_fn