[Flax] Fix mt5 auto (#12612)

Commit

4 years ago

[Flax] Fix mt5 auto (#12612) * fix_torch_device_generate_test * remove @ * fix mt5 auto