TensorFlow Recommenders: InvalidArgumentError: indices[0,1] = 66521 is not in [0, 12976) [Op:ResourceGather]

I am trying to build a recommendation system using the TensorFlow Recommenders (TFRS), and I'm facing an issue while inferencing the trained model. I have created the Query and Candidate Towers and set the embedding dimension to be the length of unique movies in my dataset. The user features have already been normailzed.

class UserModel(tf.keras.Model):
    def __init__(self, users_vocab, embedding_dimension = 128):
        super().__init__()
        self.user_model = tf.keras.Sequential([
            tf.keras.layers.experimental.preprocessing.IntegerLookup(vocabulary=users_vocab),
            tf.keras.layers.Embedding(len(users_vocab) + 2, embedding_dimension)
        ])

    def call(self, inputs):
        # Take the input dictionary, pass it through each input layer,
        # and concatenate the result.
        ret = tf.concat([
            self.user_model(inputs["USER_ID"]),
            tf.reshape(inputs["USER_FEATURE"], [-1, 1])
        ], axis=1)
        print(tf.shape(ret))
        return ret

After training when I try to get recommendations by creating index using the BruteForce Layer, and d is the dictionary containing UserId and User Features

index = tfrs.layers.factorized_top_k.BruteForce(model.query_model)
index.index(movie_features.batch(16192).map(model.candidate_model), movies)
_, titles = index(d, k=1)

My model recommends movies which are not there in my embedding vector space and throws this error when I try to get recommendations for k>1

InvalidArgumentError: indices[0,1] = 66521 is not in [0, 12976) [Op:ResourceGather]

Whereas when I pass k=1, the model gives me a movie ID

recommendations for user [2471]: [Movie ID]

How do I resolve this issue?

Thank You! :)

EDIT 1: To add to the above problem, when I pass candidate features without the identifier index.index(movie_features.batch(16192).map(model.candidate_model)) then my model returns indices correctly, but these indices values are greater than the length of movies variable.

movie_features: <MapDataset shapes: {MOVIE_ID: (), MOVIE_FEATURE: ()}, types: {MOVIE_ID: tf.int32, MOVIE_FEATURE: tf.float64}>

movies: list of unique movie id's movies = training_data.MOVIE_ID.unique() of length 83349