Masked BERT model from keras_nlp produces weird results

[vctrs]

Description

I am trying to use keras-nlp with a pretrained masked BERT model to predict some tokens in a sequence. However the model produces inconsistent results. Am I doing something wrong?

To Reproduce

https://colab.research.google.com/drive/1kL5oXveRen9Zub_7kJCI28Gg34GfDt_x

Expected behavior

Tokens in predictions should match tokens in labels. The preprocessor masks two tokens, and first two tokens of predictions probably should match labels.

[mehtamansi29]

Hi @vctrs -

Thanks for reporting the issue. I have tested the code snippet and reproduces the reported behaviour.
Attached gist file for reference.

We will look into the issue and update you the same.