Problem batching a dataset with tensorflow Dataset API

Question

I\'m trying to train a BERT model for classification with tensorflow (2.3.0).

For this, I\'m using the tensorflow Dataset API, having as inputs an input, token type i