'context' не определен при открытии модели Keras, которая использует Tensorflow Backend

Question

Я определил уровень квантования с помощью лямбды, а затем сохранил модель и не смог импортировать ее снова. Слой квантования определяет градиент, используя custom_gradient, а py_func использует внешний пакет.

save_model.py

import numpy as np
import tensorflow as tf
from keras import Input
from keras.layers import Lambda
from keras.models import Model,load_model
data = np.random.rand(10,10)

@tf.custom_gradient
def round_test(x):
    from keras import backend as K
    def grad(dy):
        return 1.0*dy
    return K.tf.py_func(lambda a: K.np.round(a,2), [x], K.tf.float32), grad

def model_(x):

    y  = Lambda(round_test, output_shape=[10])(x)
    return y

# input model
image_tensor = Input(shape=[10])
# output model
network_output = model_(image_tensor)
model_ = Model(inputs=[image_tensor], outputs=[network_output])
print(model_.summary())
model_.predict(data)

# save model
model_.save('model.h5')

load_model.py

# load model
model_s = load_model('model.h5')
model_s.predict(data)

Ошибка: NameError: имя 'context' не определено

NameError                                 Traceback (most recent call last)
<ipython-input-31-a8ddf4691f83> in <module>
      1 # model_s = load_model('model.h5', custom_objects={"bin": round_test})
----> 2 model_s = load_model('model.h5')
      3 model_s.predict(data)

~\Anaconda3\envs\py3\lib\site-packages\keras\engine\saving.py in load_model(filepath, custom_objects, compile)
    417     f = h5dict(filepath, 'r')
    418     try:
--> 419         model = _deserialize_model(f, custom_objects, compile)
    420     finally:
    421         if opened_new_file:

~\Anaconda3\envs\py3\lib\site-packages\keras\layers\core.py in decorated(*args, **kwargs)
    106   def decorated(*args, **kwargs):
    107     """Decorated function with custom gradient."""
--> 108     if context.executing_eagerly():
    109       return _eager_mode_decorator(f, *args, **kwargs)
    110     else:

NameError: name 'context' is not defined