Turned out it was a tokenization issue. The tokenizer for the model checkpoint i was using needed a do_lower_case=False flag, this was of course missing when just taking the provided script as it is. Didn’t realize that it only trained on [CLS,UNK,SEP] sequences because with TPUs I had no intuition for training speed.