Fine tuning LayoutLM: Gradients not updated error in training

I am using the Hugging face transformer for the first time. I am trying to use LayoutLM (link)
So I annotated 2 images using UBIAI tool. Then I created a hugging face dataset with column names as required by the model. Then since I am more familiar with tensorflow, I created tf.data.Dataset using the hugging face dataset model.prepare_tf_dataset and tried to train the model. I get the following error while running:

WARNING:tensorflow:Gradients do not exist for variables ['tf_layout_lm_for_token_classification_3/layoutlm/pooler/dense/kernel:0', 'tf_layout_lm_for_token_classification_3/layoutlm/pooler/dense/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument?
WARNING:tensorflow:Gradients do not exist for variables ['tf_layout_lm_for_token_classification_3/layoutlm/pooler/dense/kernel:0', 'tf_layout_lm_for_token_classification_3/layoutlm/pooler/dense/bias:0'] when minimizing the loss. If you're using `model.compile()`, did you forget to provide a `loss` argument?

I created Google Colab notebook with the full code. Anyone experienced, please I would be very thankful if you could figure out where I am going wrong. Here is the link to google collab which should end to end run except for that error.

Thank you so so much!