'Retrieval' type of code causes error InvalidArgumentError: indices[7505] = 672 is not in [0, 672)

My code is similar to the retrieval sample. I get an error as the below. Any ideas? (looked at some posts on stackoverflow, they don't seem relevant).

I don't have any dataset.cache().take(k).repeat() occurrences, simply this as in the sample:

cached_train = train.shuffle(100_000).batch(8192).cache()
cached_test = test.batch(4096).cache()

Stack:

WARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int32
WARNING:tensorflow:Gradients do not exist for variables ['counter:0'] when minimizing the loss.
WARNING:tensorflow:The dtype of the source tensor must be floating (e.g. tf.float32) when calling GradientTape.gradient, got tf.int32
WARNING:tensorflow:Gradients do not exist for variables ['counter:0'] when minimizing the loss.

Traceback (most recent call last):
File "/mnt/tmp/spark-9016e2b7-816e-4941-abde-dd6c43c753e8/recsys_tfrs_proto.py", line 352, in
main(sys.argv)
File "/mnt/tmp/spark-9016e2b7-816e-4941-abde-dd6c43c753e8/recsys_tfrs_proto.py", line 104, in main
model = create_and_train_model(movies_ds, test, train, unique_movie_titles, unique_user_ids)
File "/mnt/tmp/spark-9016e2b7-816e-4941-abde-dd6c43c753e8/recsys_tfrs_proto.py", line 181, in create_and_train_model
train_and_evaluate(cached_test, cached_train, model, 3)
File "/mnt/tmp/spark-9016e2b7-816e-4941-abde-dd6c43c753e8/recsys_tfrs_proto.py", line 191, in train_and_evaluate
model.fit(cached_train, epochs=num_epochs, verbose=0)
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 1100, in fit
tmp_logs = self.train_function(iterator)
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 828, in call
result = self._call(*args, **kwds)
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/eager/def_function.py", line 888, in _call
return self._stateless_fn(*args, **kwds)
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 2943, in call
filtered_flat_args, captured_inputs=graph_function.captured_inputs) # pylint: disable=protected-access
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 1919, in _call_flat
ctx, args, cancellation_manager=cancellation_manager))
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/eager/function.py", line 560, in call
ctx=ctx)
File "/home/hadoop/.local/lib/python3.7/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: indices[7505] = 672 is not in [0, 672)
[[node sequential/embedding/embedding_lookup (defined at mnt/tmp/spark-9016e2b7-816e-4941-abde-dd6c43c753e8/recsys_tfrs_proto.py:343) ]] [Op:__inference_train_function_530001]

Function call stack:
train_function

2021-01-07 23:35:27.709152: W tensorflow/core/kernels/data/cache_dataset_ops.cc:757] The calling iterator did not fully read the dataset being cached. In order to avoid unexpected truncation of the dataset, the partially cached contents of the dataset will be discarded. This can happen if you have an input pipeline similar to dataset.cache().take(k).repeat(). You should use dataset.take(k).cache().repeat() instead.

Start by reproducing the failure from the reported retrieval sample and inspect recsys_tfrs_proto.py, especially create_and_train_model, train_and_evaluate, and line 343 where the embedding lookup fails. Compare the input identifiers with the embedding vocabulary size shown in the error, and treat the issue as done when the cause is confirmed and a reproducible fix or clear guidance is documented.