Issue while saving MapDataset to TF index

Hello people of TF,

I'd like to ask you about the issue I get after upgrading from 0.5 to 0.6 version of tf-recommenders (wanted to upgrade all my TF stuff to 2.6.0). I do get this traceback:

    self.model.save_serving_model(self.model_management_config.exported_model_dir, self.get_entity_url_hashes())
  File "/usr/local/lib/python3.7/dist-packages/doporucovani_hpfeed_model_tf/models/retrieval.py", line 98, in save_serving_model
    index.index(entities_ds.batch(512).map(self.model.entity_model), entities)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow_recommenders/layers/factorized_top_k.py", line 519, in index
    if tf.rank(candidates) != 2:
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/util/dispatch.py", line 206, in wrapper
    return target(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/array_ops.py", line 839, in rank
    return rank_internal(input, name, optimize=True)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/ops/array_ops.py", line 859, in rank_internal
    input = ops.convert_to_tensor(input)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/profiler/trace.py", line 163, in wrapped
    return func(*args, **kwargs)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ops.py", line 1566, in convert_to_tensor
    ret = conversion_func(value, dtype=dtype, name=name, as_ref=as_ref)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py", line 346, in _constant_tensor_conversion_function
    return constant(v, dtype=dtype, name=name)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py", line 272, in constant
    allow_broadcast=True)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py", line 283, in _constant_impl
    return _constant_eager_impl(ctx, value, dtype, shape, verify_shape)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py", line 308, in _constant_eager_impl
    t = convert_to_eager_tensor(value, ctx, dtype)
  File "/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/constant_op.py", line 106, in convert_to_eager_tensor
    return ops.EagerTensor(value, ctx.device_name, dtype)
ValueError: Attempt to convert a value (<MapDataset shapes: (None, 64), types: tf.float32>) with an unsupported type (<class 'tensorflow.python.data.ops.dataset_ops.MapDataset'>) to a Tensor.
2021-09-23 10:50:29 ERROR [71] doporucovani_hpfeed_trainer_tf.task.Train Subprocess exited with code 1 {'code': 1}

Any idea how to solve this? Sorry if this is a lame question, I do consider myself as a TF noob :)

Thanks and regards,
Alan