BruteForce layer not working after loading a saved model

[Toby1218]

I have trained a deep retrieval model with NCF. Yet when saving the model, I tried both methods of saving and loading by keras and tf. I cannot specify the top_k items to retrieve from the BruteForce layer.

brute = tfrs.layers.factorized_top_k.BruteForce(model.user_model)
ds2 = items.map(lambda x: (x['parent_asin'], model.item_model(x)))
brute.index_from_dataset(ds2)
brute(tf.constant(['B1000000028']), k = 15)
brute.save('nc')

loaded = tf.keras.models.load_model('nc')
print(loaded(tf.constant(['B1000000028']), k = 5))

It returned such error

ValueError: Could not find matching concrete function to call loaded from the SavedModel. Got:
  Positional arguments (2 total):
    * <tf.Tensor 'queries:0' shape=(1,) dtype=string>
    * 5
  Keyword arguments: {'training': False}

 Expected these arguments to match one of the following 2 option(s):

Option 1:
  Positional arguments (2 total):
    * TensorSpec(shape=(None,), dtype=tf.string, name='input_1')
    * None
  Keyword arguments: {'training': True}

Option 2:
  Positional arguments (2 total):
    * TensorSpec(shape=(None,), dtype=tf.string, name='input_1')
    * None
  Keyword arguments: {'training': False}

I want to know how can I specify the number of items to retrieve after saving the model.

[4whomtbts]

@Toby1218
Hello, I'm struggling same issue.. Any progress??

[Toby1218]

@4whomtbts I got pass the error a couple months ago by using tf.saved_model.save() instead of the keras format for saving.
(Working on Tensorflow 2.15.1)

def save(self):
        brute = tfrs.layers.factorized_top_k.BruteForce(self.user_model, k = 50)
        ds = items.map(lambda x: (x['parent_asin'], self.item_model(x)))
        brute.index_from_dataset(ds)
        tf.saved_model.save(brute, 'rec_sys/model/ncf')

To load the model
self.ncf_drm = saved_model.load(f'{folder}/ncf')