Use DCN Rank example, ValueError: Could not find matching function to call loaded from the SavedModel

In my case, I use the DCN example to build the rank model. I try to persist the model and reload the model to make predictions. My code is simalar to the demo, and I add the fellow code in the run_models function (Run in Google Colab )

def run_models(use_cross_layer, deep_layer_sizes, projection_dim=None, num_runs=5):
  models = []
  rmses = []

  for i in range(num_runs):
    model = DCN(use_cross_layer=use_cross_layer,
                deep_layer_sizes=deep_layer_sizes,
                projection_dim=projection_dim)
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate))
    models.append(model)

    model.fit(cached_train, epochs=epochs, verbose=False)
    metrics = model.evaluate(cached_test, return_dict=True)
    rmses.append(metrics["RMSE"])
   /**
     *  my add code start
     */
    model.save("test-saved-model")
    reload_model = tf.keras.models.load_model("test-saved-model")
    reload_model.predict(cached_test)
   /**
     *  my add code end
     */


  mean, stdv = np.average(rmses), np.std(rmses)

  return {"model": models, "mean": mean, "stdv": stdv}

and the error info is

ValueError: in user code:

    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:1478 predict_function  *
        return step_function(self, iterator)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:1468 step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:1259 run
        return self._extended.call_for_each_replica(fn, args=args, kwargs=kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:2730 call_for_each_replica
        return self._call_for_each_replica(fn, args, kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/distribute/distribute_lib.py:3417 _call_for_each_replica
        return fn(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:1461 run_step  **
        outputs = model.predict_step(data)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/training.py:1434 predict_step
        return self(x, training=False)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/engine/base_layer.py:1012 __call__
        outputs = call_fn(inputs, *args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/saving/saved_model/utils.py:73 return_outputs_and_add_losses
        outputs, losses = fn(inputs, *args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/saving/saved_model/utils.py:171 wrap_with_training_arg
        lambda: replace_training_and_call(False))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/utils/control_flow_util.py:115 smart_cond
        pred, true_fn=true_fn, false_fn=false_fn, name=name)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/smart_cond.py:56 smart_cond
        return false_fn()
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/saving/saved_model/utils.py:171 <lambda>
        lambda: replace_training_and_call(False))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/keras/saving/saved_model/utils.py:167 replace_training_and_call
        return wrapped_call(*args, **kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py:828 __call__
        result = self._call(*args, **kwds)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py:871 _call
        self._initialize(args, kwds, add_initializers_to=initializers)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py:726 _initialize
        *args, **kwds))
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py:2969 _get_concrete_function_internal_garbage_collected
        graph_function, _ = self._maybe_define_function(args, kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py:3361 _maybe_define_function
        graph_function = self._create_graph_function(args, kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/function.py:3206 _create_graph_function
        capture_by_value=self._capture_by_value),
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py:990 func_graph_from_py_func
        func_outputs = python_func(*func_args, **func_kwargs)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/eager/def_function.py:634 wrapped_fn
        out = weak_wrapped_fn().__wrapped__(*args, **kwds)
    /usr/local/lib/python3.7/dist-packages/tensorflow/python/saved_model/function_deserialization.py:273 restored_function_body
        "\n\n".join(signature_descriptions)))

    ValueError: Could not find matching function to call loaded from the SavedModel. Got:
      Positional arguments (2 total):
        * {'movie_id': <tf.Tensor 'features_1:0' shape=(None, 1) dtype=string>, 'user_id': <tf.Tensor 'features_3:0' shape=(None, 1) dtype=string>, 'user_rating': <tf.Tensor 'features_5:0' shape=(None, 1) dtype=float32>, 'user_gender': <tf.Tensor 'features_2:0' shape=(None, 1) dtype=int32>, 'user_zip_code': <tf.Tensor 'features_6:0' shape=(None, 1) dtype=string>, 'user_occupation_text': <tf.Tensor 'features_4:0' shape=(None, 1) dtype=string>, 'bucketized_user_age': <tf.Tensor 'features:0' shape=(None, 1) dtype=int32>}
        * False
      Keyword arguments: {}
    
    Expected these arguments to match one of the following 4 option(s):
    
    Option 1:
      Positional arguments (2 total):
        * {'user_gender': TensorSpec(shape=(None,), dtype=tf.int32, name='user_gender'), 'movie_id': TensorSpec(shape=(None,), dtype=tf.string, name='movie_id'), 'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='user_id'), 'bucketized_user_age': TensorSpec(shape=(None,), dtype=tf.int32, name='bucketized_user_age'), 'user_zip_code': TensorSpec(shape=(None,), dtype=tf.string, name='user_zip_code'), 'user_occupation_text': TensorSpec(shape=(None,), dtype=tf.string, name='user_occupation_text')}
        * True
      Keyword arguments: {}
    
    Option 2:
      Positional arguments (2 total):
        * {'user_zip_code': TensorSpec(shape=(None,), dtype=tf.string, name='features/user_zip_code'), 'user_occupation_text': TensorSpec(shape=(None,), dtype=tf.string, name='features/user_occupation_text'), 'user_gender': TensorSpec(shape=(None,), dtype=tf.int32, name='features/user_gender'), 'movie_id': TensorSpec(shape=(None,), dtype=tf.string, name='features/movie_id'), 'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='features/user_id'), 'bucketized_user_age': TensorSpec(shape=(None,), dtype=tf.int32, name='features/bucketized_user_age')}
        * False
      Keyword arguments: {}
    
    Option 3:
      Positional arguments (2 total):
        * {'movie_id': TensorSpec(shape=(None,), dtype=tf.string, name='features/movie_id'), 'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='features/user_id'), 'bucketized_user_age': TensorSpec(shape=(None,), dtype=tf.int32, name='features/bucketized_user_age'), 'user_zip_code': TensorSpec(shape=(None,), dtype=tf.string, name='features/user_zip_code'), 'user_occupation_text': TensorSpec(shape=(None,), dtype=tf.string, name='features/user_occupation_text'), 'user_gender': TensorSpec(shape=(None,), dtype=tf.int32, name='features/user_gender')}
        * True
      Keyword arguments: {}
    
    Option 4:
      Positional arguments (2 total):
        * {'user_gender': TensorSpec(shape=(None,), dtype=tf.int32, name='user_gender'), 'movie_id': TensorSpec(shape=(None,), dtype=tf.string, name='movie_id'), 'user_id': TensorSpec(shape=(None,), dtype=tf.string, name='user_id'), 'bucketized_user_age': TensorSpec(shape=(None,), dtype=tf.int32, name='bucketized_user_age'), 'user_zip_code': TensorSpec(shape=(None,), dtype=tf.string, name='user_zip_code'), 'user_occupation_text': TensorSpec(shape=(None,), dtype=tf.string, name='user_occupation_text')}
        * False
      Keyword arguments: {}

In #96 , model.save_weights() and then using model.load_weights() can work.
However, in my usage scenario, I need use the function tf.keras.models.load_model
How to solve this issues? Tks again.