TFF : Modify the value of state

Question

The state object returned by iterative_process.initialize() is typically a Python container (tuple, collections.OrderedDict, etc) that contains numpy arrays. I would like that the value of state is not random, instead it begin from loaded model. As the beginning, I write this :

def create_keras_model():   
  Model = tf.keras.models.load_model(path)
  return Model
def model_fn(): 
    keras_model = create_keras_model()   
    return tff.learning.from_keras_model(keras_model..)
iterative_process = tff.learning.build_federated_averaging_process(model_fn=model_fn..)
state = iterative_process.initialize()

But test accuracy result does not change at all comparing by the normal case(if I don't load an external model).

That's why, I try this solution:

# initialize_fn() function
@tff.tf_computation
def server_init():
    model = model_fn()
    return model.trainable_variables

@tff.federated_computation
def initialize_fn():
    return tff.federated_value(server_init(), tff.SERVER) 

iterative_process = tff.templates.IterativeProcess(initialize_fn, next_fn)
state = iterative_process.initialize()
state['model'] = create_keras_model()

But I find this error:

NameError: name 'next_fn' is not defined

So in my case, how can I define next_fn ? Thanks