Why my model only predicts KEEP tags?

Hi,
I used your code to train the model in GEC task, I upgraded the code to tensorflow2.x, but when I'm in 'prediction' step, that is, when I inference the model, the model only predicts [KEEP, KEEP, ...., KEEP]. I used your original dataset, alse all KEEP in lang8-english with 20w examples, I trained 3 epochs(17w steps), only 0step and 100step model predicted not all KEEP(obviously a random output).
The same problem arises in lasertaggerff, my label map size is 2002, and follow the steps exactly, not only trained on the original few-shot dataset but also lang8, the model behaves the same, that is, predicting all KEEP with no change in src.