looks like i pasted together data batching code that doesn't line up basically the code needs to be mutated such that each batch is a dict, rather than the whole data. the example at https://github.com/huggingface/transformers/blob/master/examples/flax/language-modeling/run_t5_mlm_flax.py uses transformers.BatchEncoding for this.