MLM Using AlBert - No loss error

I’m getting this error while trying to train Albert model MLM task. (albert-xlarge-v2)

WARNING:tensorflow:Gradients do not exist for variables [‘tf_albert_for_masked_lm/predictions/bias:0’] when minimizing the loss. If you’re using model.compile(), did you forget to provide a loss argument?

The problem is not the warning but what happens next. The colab notebook just crashes and runtime is closed.

It’s not the memory issue (I am using only 10,000 lines of text for MLM which after preparation drops down to 4000 lines)