transformers

d5a72b6e - [Flax] Allow dataclasses to be jitted (#11886)

[Flax] Allow dataclasses to be jitted (#11886) * fix_torch_device_generate_test * remove @ * change dataclasses to flax ones * fix typo * fix jitted tests * fix bert & electra