DCN model prediction time increases linearly for each iteration.

Greetings to all,

I am training a DCN model for ranking purpose. The DCN model has the below structure

    tf.keras.Sequential([
      tfrs.layers.dcn.Cross(projection_dim=None,kernel_initializer="glorot_uniform"),
      tf.keras.layers.Dense(256, activation="relu"),
      tf.keras.layers.Dropout(0.25),
      tf.keras.layers.Dense(64, activation="relu"),
      tf.keras.layers.Dropout(0.10),
      tf.keras.layers.Dense(32, activation="relu"),
      tf.keras.layers.Dense(1)
  ])

The model is saved successfully and loaded in another session with the below lines of code.

tf.saved_model.save(model,model_path)
loaded_model = tf.saved_model.load(model_path)

The model takes 4 input features. So one input contains

input_data  = {'feat_0': list_of_values ,'feat_1':list_of_values, 'feat_2':list_of_values, 'feat_3':list_of_values}
prediction = loaded_model(input_data)

The issue here is that, I have more than 5000 inputs for prediction. Each with the above structure. When I iterate over these inputs and do the prediction, the time increases linearly with each iteration as shown below. But individually the prediction time is good.

Another behavior I found is that, if the model is loaded from the disc for each iteration the time remains constant. But this cannot be the solution because loading model from disc takes additional 1-2 seconds.

What could be the solution for this? Any help would be greatly appreciated. Thank you so much.